Allow Topology.add_molecule to take list[Molecule]

devtools/conda-envs/test_env.yaml

@@ -32,6 +32,7 @@ dependencies:
   - pytest-cov
   - pytest-xdist
   - pytest-rerunfailures
+  - pytest-timeout
   - pyyaml
   - toml
   - bson

openff/toolkit/_tests/test_topology.py

@@ -144,6 +144,31 @@ def test_empty(self):
         assert not topology.is_periodic
         assert len(topology.constrained_atom_pairs.items()) == 0
 
+    @pytest.mark.timeout(5)
+    def test_from_molecule_multiple(self):
+        """
+        Test that add_molecule on a list of many molecules is quick.
+
+        See issue #1916
+        """
+        water = create_water()
+
+        topology = Topology()
+
+        topology.add_molecule(10_000 * [water])
+
+        assert topology.n_molecules == 10_000
+
+    def test_from_molecule_bad_argument(self):
+        with pytest.raises(
+            ValueError,
+            match="Invalid type.*Topology",
+        ):
+
+            topology = Topology()
+
+            topology.add_molecule(create_water().to_topology())
+
     def test_reinitialization_box_vectors(self):
         topology = Topology()
         assert Topology(topology).box_vectors is None

openff/toolkit/topology/topology.py

@@ -27,6 +27,7 @@
     Optional,
     TextIO,
     Union,
+    overload,
 )
 
 import numpy as np
@@ -2390,11 +2391,39 @@ def bond(self, bond_topology_index: int) -> "Bond":  # type: ignore[return]
                 return molecule.bond(bond_molecule_index)
             this_molecule_start_index += molecule.n_bonds
 
-    def add_molecule(self, molecule: MoleculeLike) -> int:
-        """Add a copy of the molecule to the topology"""
-        idx = self._add_molecule_keep_cache(molecule)
-        self._invalidate_cached_properties()
-        return idx
+    @overload
+    def add_molecule(self, molecule: MoleculeLike) -> int: ...
+
+    @overload
+    def add_molecule(self, molecule: list[MoleculeLike]) -> list[int]: ...
+
+    def add_molecule(
+        self,
+        molecule: MoleculeLike | list[MoleculeLike],
+    ) -> int | list[int]:
+        """Add a molecule or multiple molecules to the topology."""
+        if isinstance(molecule, list):
+
+            indices = [
+                self._add_molecule_keep_cache(iter_molecule)
+                for iter_molecule in molecule
+            ]
+
+            self._invalidate_cached_properties()
+
+            return indices
+
+        elif isinstance(molecule, (Molecule, _SimpleMolecule)):
+
+            idx = self._add_molecule_keep_cache(molecule)
+
+            self._invalidate_cached_properties()
+
+            return idx
+
+        else:
+
+            raise ValueError(f"Invalid type {type(molecule)} for Topology.add_molecule")
 
     def _add_molecule_keep_cache(self, molecule: MoleculeLike) -> int:
         self._molecules.append(deepcopy(molecule))

setup.cfg

@@ -23,7 +23,7 @@ exclude_lines =
 
 [flake8]
 max-line-length = 119
-ignore = E203,W605,W503
+ignore = E203,W605,W503,E704
 exclude =
     openff/toolkit/_tests/_stale_tests.py
 per-file-ignores =