Allow accessing atoms as frame.topology.atoms

This should be equivalent to frame.atoms

chemfiles/atom.py

@@ -34,7 +34,7 @@ def __init__(self, name, type=None):
             self.type = type
 
     def __copy__(self):
-        return Atom.from_mutable_ptr(self.ffi.chfl_atom_copy(self.ptr))
+        return Atom.from_mutable_ptr(None, self.ffi.chfl_atom_copy(self.ptr))
 
     def __repr__(self):
         name = self.name
@@ -154,7 +154,7 @@ def __getitem__(self, name):
                 "Invalid type {} for an atomic property name".format(type(name))
             )
         ptr = self.ffi.chfl_atom_get_property(self.ptr, name.encode("utf8"))
-        return Property.from_mutable_ptr(ptr).get()
+        return Property.from_mutable_ptr(self, ptr).get()
 
     def __setitem__(self, name, value):
         """

chemfiles/cell.py

@@ -60,7 +60,7 @@ def __init__(self, a, b, c, alpha=90.0, beta=90.0, gamma=90.0):
         super(UnitCell, self).__init__(ptr, is_const=False)
 
     def __copy__(self):
-        return UnitCell.from_mutable_ptr(self.ffi.chfl_cell_copy(self.ptr))
+        return UnitCell.from_mutable_ptr(None, self.ffi.chfl_cell_copy(self.ptr))
 
     def __repr__(self):
         return "UnitCell({}, {}, {}, {}, {}, {})".format(*(self.lengths + self.angles))

chemfiles/frame.py

@@ -33,7 +33,7 @@ def __getitem__(self, index):
             raise IndexError("atom index ({}) out of range for this frame".format(index))
         else:
             ptr = self.frame.ffi.chfl_atom_from_frame(self.frame.mut_ptr, c_uint64(index))
-            return Atom.from_mutable_ptr(ptr)
+            return Atom.from_mutable_ptr(self, ptr)
 
     def __iter__(self):
         for i in range(len(self)):
@@ -59,7 +59,7 @@ def __init__(self):
         super(Frame, self).__init__(self.ffi.chfl_frame(), is_const=False)
 
     def __copy__(self):
-        return Frame.from_mutable_ptr(self.ffi.chfl_frame_copy(self.ptr))
+        return Frame.from_mutable_ptr(None, self.ffi.chfl_frame_copy(self.ptr))
 
     def __repr__(self):
         return "Frame with {} atoms".format(len(self.atoms))
@@ -202,7 +202,7 @@ def cell(self):
         :py:class:`Frame`. Any modification to the cell will be reflected in
         the frame.
         """
-        return UnitCell.from_mutable_ptr(self.ffi.chfl_cell_from_frame(self.mut_ptr))
+        return UnitCell.from_mutable_ptr(self, self.ffi.chfl_cell_from_frame(self.mut_ptr))
 
     @cell.setter
     def cell(self, cell):
@@ -217,7 +217,7 @@ def topology(self):
         Get read-only access to the :py:class:`Topology` of this
         :py:class:`Frame`.
         """
-        return Topology.from_const_ptr(self.ffi.chfl_topology_from_frame(self.ptr))
+        return Topology.from_const_ptr(self, self.ffi.chfl_topology_from_frame(self.ptr))
 
     @topology.setter
     def topology(self, topology):
@@ -311,7 +311,7 @@ def __getitem__(self, name):
                 "Invalid type {} for a frame property name".format(type(name))
             )
         ptr = self.ffi.chfl_frame_get_property(self.ptr, name.encode("utf8"))
-        return Property.from_mutable_ptr(ptr).get()
+        return Property.from_mutable_ptr(self, ptr).get()
 
     def __setitem__(self, name, value):
         """

chemfiles/residue.py

@@ -81,7 +81,7 @@ def __init__(self, name, resid=None):
         super(Residue, self).__init__(ptr, is_const=False)
 
     def __copy__(self):
-        return Residue.from_mutable_ptr(self.ffi.chfl_residue_copy(self.ptr))
+        return Residue.from_mutable_ptr(None, self.ffi.chfl_residue_copy(self.ptr))
 
     def __repr__(self):
         return "Residue('{}') with {} atoms".format(self.name, len(self.atoms))
@@ -119,7 +119,7 @@ def __getitem__(self, name):
                 "Invalid type {} for a residue property name".format(type(name))
             )
         ptr = self.ffi.chfl_residue_get_property(self.ptr, name.encode("utf8"))
-        return Property.from_mutable_ptr(ptr).get()
+        return Property.from_mutable_ptr(self, ptr).get()
 
     def __setitem__(self, name, value):
         """

chemfiles/selection.py

@@ -29,7 +29,7 @@ def __init__(self, selection):
         super(Selection, self).__init__(ptr, is_const=False)
 
     def __copy__(self):
-        return Selection.from_mutable_ptr(self.ffi.chfl_selection_copy(self.ptr))
+        return Selection.from_mutable_ptr(None, self.ffi.chfl_selection_copy(self.ptr))
 
     def __repr__(self):
         return "Selection('{}')".format(self.string)

chemfiles/topology.py

@@ -55,7 +55,7 @@ def __getitem__(self, index):
             raise IndexError("atom index ({}) out of range for this topology".format(index))
         else:
             ptr = self.topology.ffi.chfl_atom_from_topology(self.topology.mut_ptr, c_uint64(index))
-            return Atom.from_mutable_ptr(ptr)
+            return Atom.from_mutable_ptr(self, ptr)
 
     def __iter__(self):
         for i in range(len(self)):
@@ -106,7 +106,7 @@ def __getitem__(self, index):
             raise IndexError("residue index ({}) out of range for this topology".format(index))
         else:
             ptr = self.topology.ffi.chfl_residue_from_topology(self.topology.ptr, c_uint64(index))
-            return Residue.from_const_ptr(ptr)
+            return Residue.from_const_ptr(self, ptr)
 
     def __iter__(self):
         for i in range(len(self)):
@@ -137,14 +137,22 @@ def __init__(self):
         super(Topology, self).__init__(self.ffi.chfl_topology(), is_const=False)
 
     def __copy__(self):
-        return Topology.from_mutable_ptr(self.ffi.chfl_topology_copy(self.ptr))
+        return Topology.from_mutable_ptr(None, self.ffi.chfl_topology_copy(self.ptr))
 
     def __repr__(self):
         return "Topology with {} atoms".format(len(self.atoms))
 
     @property
     def atoms(self):
-        return TopologyAtoms(self)
+        # late import to break circular dependency
+        from .frame import FrameAtoms
+
+        # if the topology comes from a frame, allow accessing atoms as
+        # frame.topology.atoms anyway
+        if self._CxxPointer__is_const:
+            return FrameAtoms(self._CxxPointer__origin)
+        else:
+            return TopologyAtoms(self)
 
     def resize(self, count):
         """
@@ -172,7 +180,7 @@ def residue_for_atom(self, index):
             )
         ptr = self.ffi.chfl_residue_for_atom(self.ptr, c_uint64(index))
         if ptr:
-            return Residue.from_const_ptr(ptr)
+            return Residue.from_const_ptr(self, ptr)
         else:
             return None
 

chemfiles/utils.py

@@ -14,12 +14,23 @@
 
 
 class CxxPointer(object):
+    # Used to prevent adding new attributes to chemfiles objects
     __frozen = False
+    # C++ allocated pointer
     __ptr = 0
-
-    def __init__(self, ptr, is_const=True):
+    # Is the C++ pointer a const pointer? This is reauired since const and
+    # non-const pointers have the same ABI, but modifying an object through
+    # a const pointer is UB and should be prevented
+    __is_const = False
+    # None for newly allocated objects, or a CxxPointer instance for objects
+    # living inside another object (typically atoms inside a frame, or
+    # residue in a topology).
+    __origin = None
+
+    def __init__(self, ptr, is_const=True, origin=None):
         self.__ptr = ptr
         self.__is_const = is_const
+        self.__origin = origin
         self.__frozen = True
         _check_handle(ptr)
 
@@ -35,24 +46,24 @@ def __setattr__(self, key, value):
         object.__setattr__(self, key, value)
 
     @classmethod
-    def from_mutable_ptr(cls, ptr):
+    def from_mutable_ptr(cls, origin, ptr):
         """Create a new instance from a mutable pointer"""
         new = cls.__new__(cls)
-        super(cls, new).__init__(ptr, is_const=False)
+        super(cls, new).__init__(ptr, is_const=False, origin=origin)
         return new
 
     @classmethod
-    def from_const_ptr(cls, ptr):
+    def from_const_ptr(cls, origin, ptr):
         """Create a new instance from a const pointer"""
         new = cls.__new__(cls)
-        super(cls, new).__init__(ptr, is_const=True)
+        super(cls, new).__init__(ptr, is_const=True, origin=origin)
         return new
 
     @property
     def mut_ptr(self):
         """Get the **mutable** C++ pointer for this object"""
         if self.__is_const:
-            raise ChemfilesError("Trying to use a const pointer for mutable access")
+            raise ChemfilesError("Trying to use a const pointer for mutable access, this is a bug")
         else:
             return self.__ptr
 

tests/frame.py

@@ -134,6 +134,10 @@ def test_iter(self):
             self.assertEqual(atom.name, "")
         self.assertEqual(i, 2)
 
+        for i, atom in enumerate(frame.topology.atoms):
+            self.assertEqual(atom.name, "")
+        self.assertEqual(i, 2)
+
     def test_property(self):
         frame = Frame()