From 0a523939529e44d93605ce893f38f2c89f58a25a Mon Sep 17 00:00:00 2001
From: csbrasnett <c.s.brasnett@rug.nl>
Date: Mon, 8 Apr 2024 17:04:13 +0200
Subject: [PATCH] fixed resid specification and added tests

---
 vermouth/processors/annotate_mut_mod.py |  7 +++--
 vermouth/tests/test_annotate_mut_mod.py | 41 +++++++++++++++++++------
 2 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/vermouth/processors/annotate_mut_mod.py b/vermouth/processors/annotate_mut_mod.py
index e03aab67e..6108675fb 100644
--- a/vermouth/processors/annotate_mut_mod.py
+++ b/vermouth/processors/annotate_mut_mod.py
@@ -210,7 +210,6 @@ def _resiter(mod, residue_graph, resspec, library, key, molecule):
                 molecule.nodes[node_idx][key] = molecule.nodes[node_idx].get(key, []) + [mod]
     return mod_found
 
-
 def annotate_modifications(molecule, modifications, mutations):
     """
     Annotate nodes in molecule with the desired modifications and mutations
@@ -251,8 +250,10 @@ def annotate_modifications(molecule, modifications, mutations):
 
     for mutmod, key, library in associations:
         for resspec, mod in mutmod:
-            # Ie. the target residue is chain specific
-            if (resspec.get('chain') is not None) and (resspec.get('chain') == chain):
+            # Ie. the target residue is chain or residue specific
+            condition0 = ((resspec.get('chain') is not None) and (resspec.get('chain') == chain))
+            condition1 = ((resspec.get('chain') is None) and (resspec.get('resid') is not None))
+            if condition0 or condition1:
                 mod_found = _resiter(mod, residue_graph, resspec, library, key, molecule)
                 if not mod_found:
                     LOGGER.warning('Residue specified by "{}" for mutation "{}" not found. ',
diff --git a/vermouth/tests/test_annotate_mut_mod.py b/vermouth/tests/test_annotate_mut_mod.py
index 9e9bd76de..f62438be0 100644
--- a/vermouth/tests/test_annotate_mut_mod.py
+++ b/vermouth/tests/test_annotate_mut_mod.py
@@ -324,14 +324,24 @@ def test_nter_cter_modifications(node_data, edge_data, expected):
         True
     ),
     (
-            [
-                {'chain': 'A', 'resname': 'ALA', 'resid': 1},
-                {'chain': 'A', 'resname': 'ALA', 'resid': 2},
-                {'chain': 'A', 'resname': 'ALA', 'resid': 3}
-            ],
-            [(0, 1), (1, 2)],
-            [({'resname': 'GLY', 'resid': 1}, 'MET')],
-            False
+        [
+            {'chain': 'A', 'resname': 'ALA', 'resid': 1},
+            {'chain': 'A', 'resname': 'ALA', 'resid': 2},
+            {'chain': 'A', 'resname': 'ALA', 'resid': 3}
+        ],
+        [(0, 1), (1, 2)],
+        [({'resname': 'GLY', 'resid': 1}, 'MET')],
+        True
+    ),
+    (
+        [
+            {'chain': 'A', 'resname': 'ALA', 'resid': 1},
+            {'chain': 'A', 'resname': 'ALA', 'resid': 2},
+            {'chain': 'A', 'resname': 'ALA', 'resid': 3}
+        ],
+        [(0, 1), (1, 2)],
+        [({'resname': 'ALA', 'resid': 1}, 'MET')],
+        False
     ),
     (
         [
@@ -356,7 +366,20 @@ def test_nter_cter_modifications(node_data, edge_data, expected):
             {'chain': 'B', 'resname': 'ALA', 'resid': 3}
         ],
         [(0, 1), (1, 2), (3, 4), (4, 5)],
-        [({'resname': 'GLY', 'resid': '1'}, 'ALA')],
+        [({'resname': 'GLY', 'resid': 1}, 'ALA')],
+        True
+    ),
+    (
+        [
+            {'chain': 'A', 'resname': 'ALA', 'resid': 1},
+            {'chain': 'A', 'resname': 'ALA', 'resid': 2},
+            {'chain': 'A', 'resname': 'ALA', 'resid': 3},
+            {'chain': 'B', 'resname': 'ALA', 'resid': 1},
+            {'chain': 'B', 'resname': 'ALA', 'resid': 2},
+            {'chain': 'B', 'resname': 'ALA', 'resid': 3}
+        ],
+        [(0, 1), (1, 2), (3, 4), (4, 5)],
+        [({'resname': 'ALA', 'resid': 1}, 'GLY')],
         False
     ),
     (