Minimal patch for JSON deserialization of _SimpleMolecules

openff/toolkit/_tests/test_topology.py

@@ -2045,6 +2045,17 @@ def test_roundtrip(self, oleic_acid, with_conformers, n_molecules, format):
         assert next(iter(roundtrip.molecules)).n_conformers == n_conformers
 
 
+    @pytest.mark.parametrize("format", ["dict", "json"])
+    def test_roundtrip_simple_molecules(self, mixed_topology, format):
+        if format == "dict":
+            roundtrip = Topology.from_dict(mixed_topology.to_dict())
+        elif format == "json":
+            roundtrip = Topology.from_json(mixed_topology.to_json())
+
+        assert roundtrip.n_molecules == mixed_topology.n_molecules
+        assert roundtrip.n_atoms == mixed_topology.n_atoms
+
+
 @pytest.mark.parametrize(
     ("n_degrees", "num_pairs"),
     [

openff/toolkit/topology/topology.py

@@ -1302,8 +1302,18 @@ def _initialize_from_dict(self, topology_dict):
             self.box_vectors = Quantity(box_vectors_unitless, box_vectors_unit)
 
         for molecule_dict in topology_dict["molecules"]:
-            new_mol = Molecule.from_dict(molecule_dict)
+            # the serialized representation doesn't store which molecule model
+            # each of these dicts is meant to be deserialized to. Usually it'll
+            # be Molecule, so try that first. If it's supposed to be a
+            # _SimpleMolecule, it'll KeyError because Molecule has more fields.
+            try:
+                new_mol = Molecule.from_dict(molecule_dict)
+            except KeyError:
+                # masks possible other ways KeyErrors could pop up
+                new_mol = _SimpleMolecule.from_dict(molecule_dict)
+
             self._add_molecule_keep_cache(new_mol)
+
         self._invalidate_cached_properties()
 
     @staticmethod