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,6 +287,10 @@ 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)
         """
+        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)

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,6 +276,10 @@ class PDBxFile(object):
         entry : str=None
             The entry ID to assign to the CIF file
         """
+        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)
 

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,13 +44,7 @@ class TestPdbFile(unittest.TestCase):
 
     def test_WriteFile(self):
         """Write a file, read it back, and make sure it matches the original."""
-        pdb1 = PDBFile('systems/triclinic.pdb')
-        output = StringIO()
-        PDBFile.writeFile(pdb1.topology, pdb1.positions, output)
-        input = StringIO(output.getvalue())
-        pdb2 = PDBFile(input)
-        output.close();
-        input.close();
+        def compareFiles(pdb1, pdb2):
             self.assertEqual(len(pdb2.positions), 8)
             for (p1, p2) in zip(pdb1.positions, pdb2.positions):
                 self.assertVecAlmostEqual(p1, p2)
@@ -65,6 +56,26 @@ class TestPdbFile(unittest.TestCase):
                 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()
+        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."""
         with open('systems/triclinic.pdb', 'rb') as infile:

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,23 +14,14 @@ 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())
+        with tempfile.TemporaryDirectory() as tempdir:
+            filename = os.path.join(tempdir, 'temp.pdbx')
+            PDBxFile.writeFile(mol.topology, mol.positions, filename, keepIds=True)
             try:
-            pdbx = PDBxFile(input)
+                pdbx = PDBxFile(filename)
             except Exception:
                 self.fail('Parser failed to read PDBx/mmCIF file')
 
-        # Close file handles
-        output.close()
-        input.close()
-
 
     def test_Triclinic(self):
         """Test parsing a file that describes a triclinic box."""