Fixed edge case and tests

if generic mutation is given to target all residues then need to check that the resspec actually contains a chain

vermouth/processors/annotate_mut_mod.py

@@ -210,13 +210,15 @@ def annotate_modifications(molecule, modifications, mutations):
                     (mutations, 'mutation', molecule.force_field.blocks)]
 
     residue_graph = make_residue_graph(molecule)
+    # Get the name of the chain in the molecule that we're looking at
     residue = {key: residue_graph.nodes[0].get(key)
                for key in 'chain resid resname insertion_code'.split()}
     chain = residue['chain']
 
     for mutmod, key, library in associations:
         for resspec, mod in mutmod:
-            if resspec.get('chain') == chain:
+            # Ie. the target residue is chain specific
+            if (resspec.get('chain') is not None) and (resspec.get('chain') == chain):
                 mod_found = False
                 for res_idx in residue_graph:
                     if residue_matches(resspec, residue_graph, res_idx):
@@ -234,6 +236,22 @@ def annotate_modifications(molecule, modifications, mutations):
                     LOGGER.warning('Mutation "{}" not found. '
                                    'Check target resid!'
                                    ''.format(_format_resname(resspec)))
+            # If instead we're targeting all residues in the chain
+            elif resspec.get(chain) == None:
+                for res_idx in residue_graph:
+                    if residue_matches(resspec, residue_graph, res_idx):
+                        mod_found = True
+                        if mod != 'none' and mod not in library:
+                            raise NameError('{} is not known as a {} for '
+                                            'force field {}'
+                                            ''.format(mod, key, molecule.force_field.name))
+                        res = residue_graph.nodes[res_idx]
+                        LOGGER.debug('Annotating {} with {} {}',
+                                     _format_resname(res), key, mod)
+                        for node_idx in res['graph']:
+                            molecule.nodes[node_idx][key] = molecule.nodes[node_idx].get(key, []) + [mod]
+
+
 
 class AnnotateMutMod(Processor):
     """
@@ -265,7 +283,6 @@ def run_molecule(self, molecule):
         annotate_modifications(molecule, self.modifications, self.mutations)
         return molecule
     def run_system(self, system):
-        print(self.mutations)
         mols = []
         for molecule in system.molecules:
             mols.append(self.run_molecule(molecule))

vermouth/tests/test_annotate_mut_mod.py

@@ -305,18 +305,18 @@ def test_nter_cter_modifications(node_data, edge_data, expected):
 @pytest.mark.parametrize('node_data, edge_data, expected', [
     (
         [
-            {'resname': 'GLY', 'resid': 1},
-            {'resname': 'ALA', 'resid': 2},
-            {'resname': 'ALA', 'resid': 3}
+            {'chain': 'A', 'resname': 'GLY', 'resid': 1},
+            {'chain': 'A', 'resname': 'ALA', 'resid': 2},
+            {'chain': 'A', 'resname': 'ALA', 'resid': 3}
         ],
         [(0, 1), (1, 2)],
         False
     ),
     (
         [
-            {'resname': 'ALA', 'resid': 1},
-            {'resname': 'ALA', 'resid': 2},
-            {'resname': 'ALA', 'resid': 3}
+            {'chain': 'A', 'resname': 'ALA', 'resid': 1},
+            {'chain': 'A', 'resname': 'ALA', 'resid': 2},
+            {'chain': 'A', 'resname': 'ALA', 'resid': 3}
         ],
         [(0, 1), (1, 2)],
         True
@@ -328,12 +328,12 @@ def test_mod_resid_not_correct(caplog, node_data, edge_data, expected):
     mol = Molecule(force_field=ForceField(FF_UNIVERSAL_TEST))
     mol.add_nodes_from(enumerate(node_data))
     mol.add_edges_from(edge_data)
-    mutation = [({'resname': 'GLY', 'resid': 1}, 'MET')]
+    mutation = [({'resname': 'GLY', 'resid': 1, 'chain': 'A'}, 'MET')]
 
     caplog.clear()
     annotate_modifications(mol, [], mutation)
 
-    if expected:        
-        assert '"GLY1" not found.' in str(caplog.records[0].getMessage())
+    if expected:
+        assert '"A-GLY1" not found.' in str(caplog.records[0].getMessage())
     else:
         assert caplog.records == []