WIP- Support json serialization

mosdef-hub · Oct 15, 2021 · 759c727 · 759c727
1 parent 802185d
commit 759c727
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/gmso/core/dihedral.py b/gmso/core/dihedral.py
@@ -186,3 +186,4 @@ class Config:
         alias_to_fields = {
             "dihedral_types": "dihedral_types_",
         }
+        json_encoders = {IndexedSet: list}
diff --git a/gmso/formats/json.py b/gmso/formats/json.py
@@ -13,9 +13,9 @@
 from gmso.core.bond import Bond
 from gmso.core.bond_type import BondType
 from gmso.core.box import Box
-from gmso.core.dihedral import Dihedral
+from gmso.core.dihedral import Dihedral, LayeredDihedral
 from gmso.core.dihedral_type import DihedralType
-from gmso.core.improper import Improper
+from gmso.core.improper import Improper, LayeredImproper
 from gmso.core.improper_type import ImproperType
 from gmso.core.pairpotential_type import PairPotentialType
 from gmso.formats.formats_registry import loads_as, saves_as
@@ -90,7 +90,9 @@ def _to_json(top, types=False, update=True):
         Bond: json_dict["bonds"],
         Angle: json_dict["angles"],
         Dihedral: json_dict["dihedrals"],
+        LayeredDihedral: json_dict["impropers"],
         Improper: json_dict["impropers"],
+        LayeredImproper: json_dict["impropers"],
         AtomType: json_dict["atom_types"],
         BondType: json_dict["bond_types"],
         AngleType: json_dict["angle_types"],
@@ -101,20 +103,29 @@ def _to_json(top, types=False, update=True):
     for connections, exclude_attr in [
         (top._bonds, "bond_type"),
         (top._angles, "angle_type"),
-        (top._dihedrals, "dihedral_type"),
-        (top._impropers, "improper_type"),
+        (top._dihedrals, {"dihedral_type", "dihedral_types"}),
+        (top._impropers, {"improper_type", "improper_types"}),
     ]:
         for connection in connections:
             connection_dict = connection.json_dict(
                 exclude={exclude_attr, "connection_members"}
+                if not isinstance(exclude_attr, set)
+                else exclude_attr.union("connection_members")
             )
             target = targets[type(connection)]
             connection_dict["connection_members"] = [
                 top.get_index(member)
                 for member in connection.connection_members
             ]
             target.append(connection_dict)
-            connection_type = getattr(connection, exclude_attr)
+            if isinstance(exclude_attr, set):
+                for attr in exclude_attr:
+                    try:
+                        connection_type = getattr(connection, attr)
+                    except AttributeError:
+                        pass
+            else:
+                connection_type = getattr(connection, exclude_attr)
             if types and connection_type:
                 connection_dict[exclude_attr] = id(connection_type)
     if types: