writeFile() accepts a filename instead of a file object (#4162)

wrappers/python/openmm/app/pdbfile.py

@@ -6,7 +6,7 @@ Simbios, the NIH National Center for Physics-Based Simulation of
 Biological Structures at Stanford, funded under the NIH Roadmap for
 Medical Research, grant U54 GM072970. See https://simtk.org.
 
-Portions copyright (c) 2012-2021 Stanford University and the Authors.
+Portions copyright (c) 2012-2023 Stanford University and the Authors.
 Authors: Peter Eastman
 Contributors:
 
@@ -277,8 +277,8 @@ class PDBFile(object):
             The Topology defining the model to write
         positions : list
             The list of atomic positions to write
-        file : file=stdout
-            A file to write to
+        file : string or file
+            the name of the file to write.  Alternatively you can pass an open file object.
         keepIds : bool=False
             If True, keep the residue and chain IDs specified in the Topology
             rather than generating new ones.  Warning: It is up to the caller to
@@ -287,9 +287,13 @@ class PDBFile(object):
         extraParticleIdentifier : string='EP'
             String to write in the element column of the ATOM records for atoms whose element is None (extra particles)
         """
-        PDBFile.writeHeader(topology, file)
-        PDBFile.writeModel(topology, positions, file, keepIds=keepIds, extraParticleIdentifier=extraParticleIdentifier)
-        PDBFile.writeFooter(topology, file)
+        if isinstance(file, str):
+            with open(file, 'w') as output:
+                PDBFile.writeFile(topology, positions, output, keepIds, extraParticleIdentifier)
+        else:
+            PDBFile.writeHeader(topology, file)
+            PDBFile.writeModel(topology, positions, file, keepIds=keepIds, extraParticleIdentifier=extraParticleIdentifier)
+            PDBFile.writeFooter(topology, file)
 
     @staticmethod
     def writeHeader(topology, file=sys.stdout):

wrappers/python/openmm/app/pdbxfile.py

@@ -6,7 +6,7 @@ Simbios, the NIH National Center for Physics-Based Simulation of
 Biological Structures at Stanford, funded under the NIH Roadmap for
 Medical Research, grant U54 GM072970. See https://simtk.org.
 
-Portions copyright (c) 2015-2020 Stanford University and the Authors.
+Portions copyright (c) 2015-2023 Stanford University and the Authors.
 Authors: Peter Eastman
 Contributors: Jason Swails
 
@@ -266,8 +266,8 @@ class PDBxFile(object):
             The Topology defining the model to write
         positions : list
             The list of atomic positions to write
-        file : file=stdout
-            A file to write to
+        file : string or file
+            the name of the file to write.  Alternatively you can pass an open file object.
         keepIds : bool=False
             If True, keep the residue and chain IDs specified in the Topology
             rather than generating new ones.  Warning: It is up to the caller to
@@ -276,8 +276,12 @@ class PDBxFile(object):
         entry : str=None
             The entry ID to assign to the CIF file
         """
-        PDBxFile.writeHeader(topology, file, entry, keepIds)
-        PDBxFile.writeModel(topology, positions, file, keepIds=keepIds)
+        if isinstance(file, str):
+            with open(file, 'w') as output:
+                PDBxFile.writeFile(topology, positions, output, keepIds, entry)
+        else:
+            PDBxFile.writeHeader(topology, file, entry, keepIds)
+            PDBxFile.writeModel(topology, positions, file, keepIds=keepIds)
 
     @staticmethod
     def writeHeader(topology, file=sys.stdout, entry=None, keepIds=False):

wrappers/python/tests/TestPdbFile.py

-import sys
+import tempfile
 import unittest
 from openmm.app import *
 from openmm import *
 from openmm.unit import *
 import openmm.app.element as elem
-if sys.version_info >= (3, 0):
-    from io import StringIO
-else:
-    from cStringIO import StringIO
+from io import StringIO
 
 
 class TestPdbFile(unittest.TestCase):
@@ -47,23 +44,37 @@ class TestPdbFile(unittest.TestCase):
 
     def test_WriteFile(self):
         """Write a file, read it back, and make sure it matches the original."""
+        def compareFiles(pdb1, pdb2):
+            self.assertEqual(len(pdb2.positions), 8)
+            for (p1, p2) in zip(pdb1.positions, pdb2.positions):
+                self.assertVecAlmostEqual(p1, p2)
+            for (v1, v2) in zip(pdb1.topology.getPeriodicBoxVectors(), pdb2.topology.getPeriodicBoxVectors()):
+                self.assertVecAlmostEqual(v1, v2, 1e-4)
+            self.assertVecAlmostEqual(pdb1.topology.getUnitCellDimensions(), pdb2.topology.getUnitCellDimensions(), 1e-4)
+            for atom1, atom2 in zip(pdb1.topology.atoms(), pdb2.topology.atoms()):
+                self.assertEqual(atom1.element, atom2.element)
+                self.assertEqual(atom1.name, atom2.name)
+                self.assertEqual(atom1.residue.name, atom2.residue.name)
+
         pdb1 = PDBFile('systems/triclinic.pdb')
+
+        # First try writing to an open file object.
+
         output = StringIO()
         PDBFile.writeFile(pdb1.topology, pdb1.positions, output)
         input = StringIO(output.getvalue())
         pdb2 = PDBFile(input)
-        output.close();
-        input.close();
-        self.assertEqual(len(pdb2.positions), 8)
-        for (p1, p2) in zip(pdb1.positions, pdb2.positions):
-            self.assertVecAlmostEqual(p1, p2)
-        for (v1, v2) in zip(pdb1.topology.getPeriodicBoxVectors(), pdb2.topology.getPeriodicBoxVectors()):
-            self.assertVecAlmostEqual(v1, v2, 1e-4)
-        self.assertVecAlmostEqual(pdb1.topology.getUnitCellDimensions(), pdb2.topology.getUnitCellDimensions(), 1e-4)
-        for atom1, atom2 in zip(pdb1.topology.atoms(), pdb2.topology.atoms()):
-            self.assertEqual(atom1.element, atom2.element)
-            self.assertEqual(atom1.name, atom2.name)
-            self.assertEqual(atom1.residue.name, atom2.residue.name)
+        output.close()
+        input.close()
+        compareFiles(pdb1, pdb2)
+
+        # Now try a filename.
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            filename = os.path.join(tempdir, 'temp.pdb')
+            PDBFile.writeFile(pdb1.topology, pdb1.positions, filename)
+            pdb2 = PDBFile(filename)
+            compareFiles(pdb1, pdb2)
 
     def test_BinaryStream(self):
         """Test reading a stream that was opened in binary mode."""

wrappers/python/tests/TestPdbxFile.py

+import tempfile
 import unittest
 from openmm.app import *
 from openmm import *
 from openmm.unit import *
 import openmm.app.element as elem
 import os
-if sys.version_info >= (3, 0):
-    from io import StringIO
-else:
-    from cStringIO import StringIO
+from io import StringIO
 
 class TestPdbxFile(unittest.TestCase):
     """Test the PDBx/mmCIF file parser"""
@@ -16,22 +14,13 @@ class TestPdbxFile(unittest.TestCase):
         """Test conversion from PDB to PDBx"""
 
         mol = PDBFile('systems/ala_ala_ala.pdb')
-
-        # Write to 'file'
-        output = StringIO()
-        PDBxFile.writeFile(mol.topology, mol.positions, output,
-                           keepIds=True)
-
-        # Read from 'file'
-        input = StringIO(output.getvalue())
-        try:
-            pdbx = PDBxFile(input)
-        except Exception:
-            self.fail('Parser failed to read PDBx/mmCIF file')
-
-        # Close file handles
-        output.close()
-        input.close()
+        with tempfile.TemporaryDirectory() as tempdir:
+            filename = os.path.join(tempdir, 'temp.pdbx')
+            PDBxFile.writeFile(mol.topology, mol.positions, filename, keepIds=True)
+            try:
+                pdbx = PDBxFile(filename)
+            except Exception:
+                self.fail('Parser failed to read PDBx/mmCIF file')
 
 
     def test_Triclinic(self):