Issue 170 lam validators

gufe/mapping/ligandatommapping.py

@@ -1,5 +1,7 @@
 # This code is part of gufe and is licensed under the MIT license.
 # For details, see https://github.com/OpenFreeEnergy/gufe
+from __future__ import annotations
+
 import json
 from typing import Any, Optional
 import numpy as np
@@ -18,6 +20,8 @@ class LigandAtomMapping(AtomMapping):
     """
     componentA: SmallMoleculeComponent
     componentB: SmallMoleculeComponent
+    _annotations: dict[str, Any]
+    _compA_to_compB: dict[int, int]
 
     def __init__(
         self,
@@ -34,18 +38,31 @@ def __init__(
         componentA_to_componentB : dict[int, int]
           correspondence of indices of atoms between the two ligands; the
           keys are indices in componentA and the values are indices in
-          componentB
+          componentB.
+          These are checked that they are within the possible indices of the
+          respective components.
         annotations : dict[str, Any]
           Mapping of annotation identifier to annotation data. Annotations may
           contain arbitrary JSON-serializable data. Annotation identifiers
           starting with ``ofe-`` may have special meaning in other parts of
           OpenFE. ``score`` is a reserved annotation identifier.
         """
         super().__init__(componentA, componentB)
+
+        # validate compA_to_compB
+        nA = self.componentA.to_rdkit().GetNumAtoms()
+        nB = self.componentB.to_rdkit().GetNumAtoms()
+        for i, j in componentA_to_componentB.items():
+            if not (0 <= i < nA):
+                raise ValueError(f"Got invalid index for ComponentA ({i}); "
+                                 f"must be 0 <= n < {nA}")
+            if not (0 <= j < nB):
+                raise ValueError(f"Got invalid index for ComponentB ({i}); "
+                                 f"must be 0 <= n < {nB}")
+
         self._compA_to_compB = componentA_to_componentB
 
         if annotations is None:
-            # TODO: this should be a frozen dict
             annotations = {}
 
         self._annotations = annotations
@@ -74,6 +91,7 @@ def componentB_unique(self):
 
     @property
     def annotations(self):
+        """Any extra metadata, for example the score of a mapping"""
         # return a copy (including copy of nested)
         return json.loads(json.dumps(self._annotations))
 
@@ -152,8 +170,8 @@ def draw_to_file(self, fname: str, d2d=None):
             f.write(draw_mapping(self._compA_to_compB, self.componentA.to_rdkit(),
                                  self.componentB.to_rdkit(), d2d))
 
-    def with_annotations(self, annotations: dict[str, Any]):
-        """Create an new mapping based on this one with extra annotations.
+    def with_annotations(self, annotations: dict[str, Any]) -> LigandAtomMapping:
+        """Create a new mapping based on this one with extra annotations.
 
         Parameters
         ----------

gufe/tests/test_ligandatommapping.py

@@ -13,11 +13,11 @@
 from .test_tokenization import GufeTokenizableTestsMixin
 
 
-def mol_from_smiles(smiles: str) -> Chem.Mol:
-    m = Chem.MolFromSmiles(smiles)
+def mol_from_smiles(smiles: str) -> gufe.SmallMoleculeComponent:
+    m = Chem.AddHs(Chem.MolFromSmiles(smiles))
     m.Compute2DCoords()
 
-    return m
+    return gufe.SmallMoleculeComponent(m)
 
 
 @pytest.fixture(scope='session')
@@ -28,8 +28,8 @@ def simple_mapping():
 
     C C
     """
-    molA = gufe.SmallMoleculeComponent(mol_from_smiles('CCO'))
-    molB = gufe.SmallMoleculeComponent(mol_from_smiles('CC'))
+    molA = mol_from_smiles('CCO')
+    molB = mol_from_smiles('CC')
 
     m = LigandAtomMapping(molA, molB, componentA_to_componentB={0: 0, 1: 1})
 
@@ -44,8 +44,8 @@ def other_mapping():
 
     C   C
     """
-    molA = SmallMoleculeComponent(mol_from_smiles('CCO'))
-    molB = SmallMoleculeComponent(mol_from_smiles('CC'))
+    molA = mol_from_smiles('CCO')
+    molB = mol_from_smiles('CC')
 
     m = LigandAtomMapping(molA, molB, componentA_to_componentB={0: 0, 2: 1})
 
@@ -236,10 +236,46 @@ def test_with_fancy_annotations(simple_mapping):
     assert m == m2
 
 
+class TestLigandAtomMappingBoundsChecks:
+    @pytest.fixture
+    def molA(self):
+        # 9 atoms
+        return mol_from_smiles('CCO')
+
+    @pytest.fixture
+    def molB(self):
+        # 11 atoms
+        return mol_from_smiles('CCC')
+
+    def test_too_large_A(self, molA, molB):
+        with pytest.raises(ValueError, match="invalid index for ComponentA"):
+            LigandAtomMapping(componentA=molA,
+                              componentB=molB,
+                              componentA_to_componentB={9: 5})
+
+    def test_too_small_A(self, molA, molB):
+        with pytest.raises(ValueError, match="invalid index for ComponentA"):
+            LigandAtomMapping(componentA=molA,
+                              componentB=molB,
+                              componentA_to_componentB={-2: 5})
+
+    def test_too_large_B(self, molA, molB):
+        with pytest.raises(ValueError, match="invalid index for ComponentB"):
+            LigandAtomMapping(componentA=molA,
+                              componentB=molB,
+                              componentA_to_componentB={5: 11})
+
+    def test_too_small_B(self, molA, molB):
+        with pytest.raises(ValueError, match="invalid index for ComponentB"):
+            LigandAtomMapping(componentA=molA,
+                              componentB=molB,
+                              componentA_to_componentB={5: -1})
+
+
 class TestLigandAtomMapping(GufeTokenizableTestsMixin):
     cls = LigandAtomMapping
     repr = "LigandAtomMapping(componentA=SmallMoleculeComponent(name=), componentB=SmallMoleculeComponent(name=), componentA_to_componentB={0: 0, 1: 1}, annotations={'foo': 'bar'})"
-    key = "LigandAtomMapping-c95a2c15fe21f446cf731338427137ae"
+    key = "LigandAtomMapping-c333723fbbee702c641cb9dca9beae49"
 
     @pytest.fixture
     def instance(self, annotated_simple_mapping):