Add the object variable to the atom model of cmd.get_model()

schrodinger · Aug 6, 2024 · 1de7c8b · 1de7c8b
1 parent 03d7a7f
commit 1de7c8b
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 15 deletions.
diff --git a/layer2/CoordSet.cpp b/layer2/CoordSet.cpp
@@ -1,16 +1,16 @@
 
-/* 
+/*
 A* -------------------------------------------------------------------
 B* This file contains source code for the PyMOL computer program
-C* copyright 1998-2000 by Warren Lyford Delano of DeLano Scientific. 
+C* copyright 1998-2000 by Warren Lyford Delano of DeLano Scientific.
 D* -------------------------------------------------------------------
 E* It is unlawful to modify or remove this copyright notice.
 F* -------------------------------------------------------------------
-G* Please see the accompanying LICENSE file for further information. 
+G* Please see the accompanying LICENSE file for further information.
 H* -------------------------------------------------------------------
 I* Additional authors of this source file include:
--* 
--* 
+-*
+-*
 -*
 Z* -------------------------------------------------------------------
 */
@@ -1054,8 +1054,8 @@ void CoordSetAtomToPDBStrVLA(PyMOLGlobals * G, char **charVLA, int *c,
 
 
 /*========================================================================*/
-PyObject *CoordSetAtomToChemPyAtom(PyMOLGlobals * G, AtomInfoType * ai, const float *v,
-                                   const float *ref, int index, const double *matrix)
+PyObject *CoordSetAtomToChemPyAtom(PyMOLGlobals * G, AtomInfoType * ai, ObjectMolecule *obj,
+                                   const float *v, const float *ref, int index, const double *matrix)
 {
 #ifdef _PYMOL_NOPY
   return nullptr;
@@ -1130,6 +1130,8 @@ PyObject *CoordSetAtomToChemPyAtom(PyMOLGlobals * G, AtomInfoType * ai, const fl
     PConvIntToPyObjAttr(atom, "id", ai->id);    /* not necc. unique */
     PConvIntToPyObjAttr(atom, "index", index + 1);      /* fragile */
 
+    PConvStringToPyObjAttr(atom, "object", obj->Name);
+
 #ifdef _PYMOL_IP_PROPERTIES
 #endif
 

diff --git a/layer2/CoordSet.h b/layer2/CoordSet.h
@@ -1,16 +1,16 @@
 
-/* 
+/*
 A* -------------------------------------------------------------------
 B* This file contains source code for the PyMOL computer program
-C* copyright 1998-2000 by Warren Lyford Delano of DeLano Scientific. 
+C* copyright 1998-2000 by Warren Lyford Delano of DeLano Scientific.
 D* -------------------------------------------------------------------
 E* It is unlawful to modify or remove this copyright notice.
 F* -------------------------------------------------------------------
-G* Please see the accompanying LICENSE file for further information. 
+G* Please see the accompanying LICENSE file for further information.
 H* -------------------------------------------------------------------
 I* Additional authors of this source file include:
--* 
--* 
+-*
+-*
 -*
 Z* -------------------------------------------------------------------
 */
@@ -122,7 +122,7 @@ struct CoordSet : CObjectState {
 
   pymol::vla<RefPosType> RefPos;
 
-  /* idea:  
+  /* idea:
      int start_atix, stop_atix <-- for discrete objects, we need
      something like this that would enable pymol to skip atoms not in the
      discrete state...question is: are these atoms sorted together right
@@ -202,8 +202,8 @@ bool CoordSetInsureOrthogonal(PyMOLGlobals * G,
     bool quiet=true);
 
 void CoordSetGetAverage(const CoordSet * I, float *v0);
-PyObject *CoordSetAtomToChemPyAtom(PyMOLGlobals * G, AtomInfoType * ai, const float *v,
-                                   const float *ref, int index, const double *matrix);
+PyObject *CoordSetAtomToChemPyAtom(PyMOLGlobals * G, AtomInfoType * ai, ObjectMolecule * obj,
+                                   const float *v, const float *ref, int index, const double *matrix);
 int CoordSetGetAtomVertex(const CoordSet * I, int at, float *v);
 int CoordSetGetAtomTxfVertex(const CoordSet * I, int at, float *v);
 int CoordSetSetAtomVertex(CoordSet * I, int at, const float *v);

diff --git a/testing/tests/api/exporting.py b/testing/tests/api/exporting.py
@@ -348,6 +348,15 @@ def testPSEBulkExportImport(self, pse_export_version, pse_binary_dump):
             m2 = cmd.get_model()
             self.assertModelsAreSame(m1, m2)
 
+    def testGetModelObjectName(self):
+        cmd.load(self.datafile("1oky-frag.pdb"))
+        cmd.load(self.datafile('1rna.cif'))
+
+        m1 = cmd.get_model()
+        self.assertEqual(m1.atom[0].object, '1oky-frag')
+        self.assertEqual(m1.atom[-1].object, '1rna')
+
+
     @testing.requires_version('2.1')
     def testMMTF(self):
         '''Styled MMTF export/import'''