Add ID checks when saving

atomium/structures.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import rmsd
+import warnings
 from collections import Counter, OrderedDict
 from .base import StructureClass, query, StructureSet
 
@@ -299,8 +300,30 @@ def equivalent_to(self, other):
         except: return False
 
 
+    def check_ids(self):
+        """Looks through all the structure's sub-structures and raises a
+        warning if they have duplicate ID."""
+
+        for objects in ("chains", "ligands", "waters", "residues", "atoms"):
+            try:
+                ids = [obj.id for obj in getattr(self, objects)()]
+                unique_ids = set(ids)
+                if len(ids) != len(unique_ids):
+                    warnings.warn(f"{objects} have duplicate IDs")
+            except AttributeError: pass
+
+
     def save(self, path):
+        """Saves the structure to file. The file extension given in the filename
+        will be used to determine which file format to save in.
+
+        If the structure you are saving has any duplicate IDs, a warning will be
+        issued, as the file saved will likely be nonsensical.
+
+        :param str path: the filename and location to save to."""
+
         from .utilities import save
+        self.check_ids()
         ext = path.split(".")[-1]
         if ext == "cif":
             from .mmcif import structure_to_mmcif_string

tests/integration/test_file_saving.py

@@ -67,6 +67,13 @@ def test_chain(self):
         self.assertTrue(f.model.chain("A").equivalent_to(chain))
 
 
+    def test_biological_assembly_warns_on_saving(self):
+        f = atomium.open("tests/integration/files/1xda.cif")
+        model = f.generate_assembly(5)
+        with self.assertWarns(Warning):
+            model.save("tests/integration/files/assembly.cif")
+
+
 
 class MmtfFileSavingTests(SavingTest):
 
@@ -101,6 +108,13 @@ def test_chain(self):
         self.assertTrue(f.model.chain("A").equivalent_to(chain))
 
 
+    def test_biological_assembly_warns_on_saving(self):
+        f = atomium.open("tests/integration/files/1xda.cif")
+        model = f.generate_assembly(5)
+        with self.assertWarns(Warning):
+            model.save("tests/integration/files/assembly.cif")
+
+
 
 class PdbFileSavingTests(SavingTest):
 
@@ -133,3 +147,10 @@ def test_chain(self):
         f.model.chain("A").save("tests/integration/files/chaina.pdb")
         chain = atomium.open("tests/integration/files/chaina.pdb").model
         self.assertTrue(f.model.chain("A").equivalent_to(chain))
+
+
+    def test_biological_assembly_warns_on_saving(self):
+        f = atomium.open("tests/integration/files/1xda.cif")
+        model = f.generate_assembly(5)
+        with self.assertWarns(Warning):
+            model.save("tests/integration/files/assembly.cif")

tests/unit/test_atom_structures.py

@@ -354,19 +354,83 @@ def test_structures_not_equivalent_if_atoms_not_equivalent(self, mock_pair):
 
 
 
+class AtomStructureIdCheckingTests(AtomStructureTest):
+
+    def test_can_warn_on_atom_ids(self):
+        self.structure.atoms = lambda : [Mock(id=1), Mock(id=2)]
+        self.structure.check_ids()
+        self.structure.atoms = lambda : [Mock(id=1), Mock(id=1)]
+        with self.assertWarns(Warning):
+            self.structure.check_ids()
+
+
+    def test_can_warn_on_residue_ids(self):
+        self.structure.residues = lambda : [Mock(id=1), Mock(id=2)]
+        self.structure.check_ids()
+        self.structure.residues = lambda : [Mock(id=1), Mock(id=1)]
+        with self.assertWarns(Warning):
+            self.structure.check_ids()
+
+
+    def test_can_warn_on_chain_ids(self):
+        self.structure.chains = lambda : [Mock(id=1), Mock(id=2)]
+        self.structure.check_ids()
+        self.structure.chains = lambda : [Mock(id=1), Mock(id=1)]
+        with self.assertWarns(Warning):
+            self.structure.check_ids()
+
+
+    def test_can_warn_on_ligand_ids(self):
+        self.structure.ligands = lambda : [Mock(id=1), Mock(id=2)]
+        self.structure.check_ids()
+        self.structure.ligands = lambda : [Mock(id=1), Mock(id=1)]
+        with self.assertWarns(Warning):
+            self.structure.check_ids()
+
+
+    def test_can_warn_on_water_ids(self):
+        self.structure.waters = lambda : [Mock(id=1), Mock(id=2)]
+        self.structure.check_ids()
+        self.structure.waters = lambda : [Mock(id=1), Mock(id=1)]
+        with self.assertWarns(Warning):
+            self.structure.check_ids()
+
+
+
 class AtomStructureSavingTests(AtomStructureTest):
 
+    @patch("atomium.structures.AtomStructure.check_ids")
     @patch("atomium.mmcif.structure_to_mmcif_string")
     @patch("atomium.utilities.save")
-    def test_can_save_cif(self, mock_save, mock_conv):
+    def test_can_save_cif(self, mock_save, mock_conv, mock_check):
         self.structure.save("test.cif")
+        mock_check.assert_called_with()
         mock_conv.assert_called_with(self.structure)
         mock_save.assert_called_with(mock_conv.return_value, "test.cif")
 
 
+    @patch("atomium.structures.AtomStructure.check_ids")
     @patch("atomium.mmtf.structure_to_mmtf_string")
     @patch("atomium.utilities.save")
-    def test_can_save_cif(self, mock_save, mock_conv):
+    def test_can_save_mmtf(self, mock_save, mock_conv, mock_check):
         self.structure.save("test.mmtf")
+        mock_check.assert_called_with()
         mock_conv.assert_called_with(self.structure)
         mock_save.assert_called_with(mock_conv.return_value, "test.mmtf")
+
+
+    @patch("atomium.structures.AtomStructure.check_ids")
+    @patch("atomium.pdb.structure_to_pdb_string")
+    @patch("atomium.utilities.save")
+    def test_can_save_pdb(self, mock_save, mock_conv, mock_check):
+        self.structure.save("test.pdb")
+        mock_check.assert_called_with()
+        mock_conv.assert_called_with(self.structure)
+        mock_save.assert_called_with(mock_conv.return_value, "test.pdb")
+
+
+    @patch("atomium.structures.AtomStructure.check_ids")
+    def test_can_reject_weird_extensions(self, mock_check):
+        with self.assertRaises(ValueError):
+            self.structure.save("test.abc")
+        mock_check.assert_called_with()