Refactor ChemfilesWriter to follow the same code as the XYZ writer

MDAnalysis · Apr 8, 2021 · 3a62f46 · 3a62f46
1 parent b3b869d
commit 3a62f46
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 7 deletions.
diff --git a/package/MDAnalysis/coordinates/chemfiles.py b/package/MDAnalysis/coordinates/chemfiles.py
@@ -342,7 +342,7 @@ def _write_next_frame(self, obj):
         constructor.
 
         If `obj` contains velocities, and the underlying format supports it, the
-        velocities are writen to the file. Writing forces is unsupported at the
+        velocities are written to the file. Writing forces is unsupported at the
         moment.
 
         Parameters
@@ -356,21 +356,23 @@ def _write_next_frame(self, obj):
            Deprecated support for Timestep argument has now been removed.
            Use AtomGroup or Universe as an input instead.
         """
-        if hasattr(obj, "atoms"):
+        try:
+            atoms = obj.atoms
+        except AttributeError:
+            errmsg = "Input obj is neither an AtomGroup or Universe"
+            raise TypeError(errmsg) from None
+        else:
             if hasattr(obj, "universe"):
                 # For AtomGroup and children (Residue, ResidueGroup, Segment)
                 ts_full = obj.universe.trajectory.ts
-                if ts_full.n_atoms == obj.atoms.n_atoms:
+                if ts_full.n_atoms == atoms.n_atoms:
                     ts = ts_full
                 else:
                     # Only populate a time step with the selected atoms.
-                    ts = ts_full.copy_slice(obj.atoms.indices)
+                    ts = ts_full.copy_slice(atoms.indices)
             elif hasattr(obj, "trajectory"):
                 # For Universe only --- get everything
                 ts = obj.trajectory.ts
-        else:
-            errmsg = "Input obj is neither an AtomGroup or Universe"
-            raise TypeError(errmsg) from None
 
         frame = self._timestep_to_chemfiles(ts)
         frame.topology = self._topology_to_chemfiles(obj, len(frame.atoms))

diff --git a/testsuite/MDAnalysisTests/coordinates/test_chemfiles.py b/testsuite/MDAnalysisTests/coordinates/test_chemfiles.py
@@ -195,6 +195,17 @@ def test_write_topology(self, tmpdir):
 
             # self.check_topology(datafiles.CONECT, outfile)
 
+    def test_write_atom_group(self, tmpdir):
+        u = mda.Universe(datafiles.CONECT)
+        group = u.select_atoms("resname ARG")
+        with tmpdir.as_cwd():
+            outfile = "chemfiles-write-atom-group.pdb"
+            with ChemfilesWriter(outfile) as writer:
+                writer.write(group)
+
+            check = mda.Universe(outfile)
+            assert check.trajectory.ts.n_atoms == group.n_atoms
+
     def test_write_velocities(self, tmpdir):
         u = mda.Universe.empty(4, trajectory=True)
         u.add_TopologyAttr("type", values=["H", "H", "H", "H"])