[feat] added tests for multiple models

Author: kierandidi

biopandas/mmcif/pandas_mmcif.py

@@ -70,6 +70,15 @@ def read_mmcif(self, path):
         self.code = self.data["entry"]["id"][0].lower()
         return self
 
+    def label_models(self):
+        """Adds a column ("model_id") to the underlying
+        DataFrames containing the model number."""
+        if "ATOM" in self.df.keys():
+            self.df["ATOM"]["model_id"] = self.df["ATOM"]["pdbx_PDB_model_num"]
+        if "HETATM" in self.df.keys():
+            self.df["HETATM"]["model_id"] = self.df["HETATM"]["pdbx_PDB_model_num"]
+        return self
+    
     def get_model(self, model_index: int) -> PandasMmcif:
         """Returns a new PandasMmcif object with the dataframes subset to the
         given model index.

biopandas/mmcif/tests/test_multiple_models.py

@@ -10,6 +10,11 @@
 
 TESTDATA_FILENAME = os.path.join(os.path.dirname(__file__), "data", "2jyf.cif.gz")
 
+def test_label_models():
+    df = PandasMmcif().read_mmcif(TESTDATA_FILENAME)
+    df.label_models()
+    assert "model_id" in df.df["ATOM"].columns
+    
 def test_get_model():
     df = PandasMmcif().read_mmcif(TESTDATA_FILENAME)
     MODEL_INDEX = 1