Fix DATAWriter, CRD, PQR, and PDBQT compressed file writing

package/CHANGELOG

@@ -19,6 +19,8 @@ The rules for this file:
 
 Fixes
   * Fix deletion of bond/angle/dihedral types (PR #4003, Issue #4001)
+  * CRD, PQR, and PDBQT writers now write BZ2 and GZ files if requested
+    (PR #4163, Issue #4159)
 
 Enhancements
 

package/MDAnalysis/coordinates/CRD.py

@@ -128,6 +128,9 @@ class CRDWriter(base.WriterBase):
     .. versionchanged:: 2.2.0
        CRD extended format can now be explicitly requested with the
        `extended` keyword
+    .. versionchanged:: 2.6.0
+       Files are now writen in `wt` mode, and keep extensions, allowing
+       for files to be written under compressed formats
     """
     format = 'CRD'
     units = {'time': None, 'length': 'Angstrom'}
@@ -165,7 +168,7 @@ def __init__(self, filename, **kwargs):
              .. versionadded:: 2.2.0
         """
 
-        self.filename = util.filename(filename, ext='crd')
+        self.filename = util.filename(filename, ext='crd', keep=True)
         self.crd = None
 
         # account for explicit crd format, if requested
@@ -247,7 +250,7 @@ def write(self, selection, frame=None):
                 "{miss}. These will be written with default values. "
                 "".format(miss=', '.join(missing_topology)))
 
-        with util.openany(self.filename, 'w') as crd:
+        with util.openany(self.filename, 'wt') as crd:
             # Write Title
             crd.write(self.fmt['TITLE'].format(
                 frame=frame, where=u.trajectory.filename))

package/MDAnalysis/coordinates/LAMMPS.py

@@ -263,7 +263,7 @@ def __init__(self, filename, convert_units=True, **kwargs):
         convert_units : bool, optional
             units are converted to the MDAnalysis base format; [``True``]
         """
-        self.filename = util.filename(filename, ext='data')
+        self.filename = util.filename(filename, ext='data', keep=True)
 
         self.convert_units = convert_units
 
@@ -419,7 +419,7 @@ def write(self, selection, frame=None):
             has_velocities = True
 
         features = {}
-        with util.openany(self.filename, 'w') as self.f:
+        with util.openany(self.filename, 'wt') as self.f:
             self.f.write('LAMMPS data file via MDAnalysis\n')
             self.f.write('\n')
             self.f.write('{:>12d}  atoms\n'.format(len(atoms)))

package/MDAnalysis/coordinates/NAMDBIN.py

@@ -94,7 +94,12 @@ def Writer(self, filename, **kwargs):
 class NAMDBINWriter(base.WriterBase):
     """Writer for NAMD binary coordinate files.
 
-
+
+    Note
+    ----
+    * Does not handle writing to bz2 or gz compressed file types.
+
+
     .. versionadded:: 1.0.0
     """
     format = ['COOR', 'NAMDBIN']

package/MDAnalysis/coordinates/PDBQT.py

@@ -195,6 +195,11 @@ class PDBQTWriter(base.WriterBase):
 
     .. _PDB: http://www.wwpdb.org/documentation/file-format-content/format32/v3.2.html
     .. _PDBQT: http://autodock.scripps.edu/faqs-help/faq/what-is-the-format-of-a-pdbqt-file
+
+
+    .. versionchanged:: 2.6.0
+       Files are now writen in `wt` mode, and keep extensions, allowing
+       for files to be written under compressed formats
     """
 
     fmt = {
@@ -213,8 +218,8 @@ class PDBQTWriter(base.WriterBase):
     pdb_coor_limits = {"min": -999.9995, "max": 9999.9995}
 
     def __init__(self, filename, **kwargs):
-        self.filename = util.filename(filename, ext='pdbqt')
-        self.pdb = util.anyopen(self.filename, 'w')
+        self.filename = util.filename(filename, ext='pdbqt', keep=True)
+        self.pdb = util.anyopen(self.filename, 'wt')
 
     def close(self):
         self.pdb.close()

package/MDAnalysis/coordinates/PQR.py

@@ -186,7 +186,11 @@ class PQRWriter(base.WriterBase):
 
     The output format is similar to ``pdb2pqr --whitespace``.
 
+
     .. versionadded:: 0.9.0
+    .. versionchanged:: 2.6.0
+       Files are now writen in `wt` mode, and keep extensions, allowing
+       for files to be written under compressed formats
     """
     format = 'PQR'
     units = {'time': None, 'length': 'Angstrom'}
@@ -210,7 +214,7 @@ def __init__(self, filename, convert_units=True, **kwargs):
              remark lines (list of strings) or single string to be added to the
              top of the PQR file
         """
-        self.filename = util.filename(filename, ext='pqr')
+        self.filename = util.filename(filename, ext='pqr', keep=True)
         self.convert_units = convert_units  # convert length and time to base units
         self.remarks = kwargs.pop('remarks', "PQR file written by MDAnalysis")
 
@@ -299,7 +303,7 @@ def write(self, selection, frame=None):
                 "{miss}. These will be written with default values. "
                 "".format(miss=', '.join(missing_topology)))
 
-        with util.openany(self.filename, 'w') as pqrfile:
+        with util.openany(self.filename, 'wt') as pqrfile:
             # Header / Remarks
             # The *remarknumber* is typically 1 but :program:`pdb2pgr`
             # also uses 6 for the total charge and 5 for warnings.

package/MDAnalysis/lib/util.py

@@ -338,9 +338,10 @@ def anyopen(datasource, mode='rt', reset=True):
         a file (from :class:`file` or :func:`open`) or a stream (e.g. from
         :func:`urllib2.urlopen` or :class:`io.StringIO`)
     mode: {'r', 'w', 'a'} (optional)
-        Open in r(ead), w(rite) or a(ppen) mode. More complicated
-        modes ('r+', 'w+', ...) are not supported; only the first letter of
-        `mode` is used and thus any additional modifiers are silently ignored.
+        Open in r(ead), w(rite) or a(ppen) mode. This string is directly
+        passed to the file opening handler (either Python's openfe, bz2, or
+        gzip). More complex mode are supported if the file opening handler
+        supports it.
     reset: bool (optional)
         try to read (`mode` 'r') the stream from the start
 

testsuite/MDAnalysisTests/coordinates/test_crd.py

@@ -43,14 +43,17 @@ def u(self):
     def outfile(self, tmpdir):
         return os.path.join(str(tmpdir), 'test.crd')
 
-    def test_write_atoms(self, u, outfile):
+    @pytest.mark.parametrize('testfile',
+        ['test.crd', 'test.crd.bz2', 'test.crd.gz'])
+    def test_write_atoms(self, u, testfile, tmpdir):
         # Test that written file when read gives same coordinates
-        u.atoms.write(outfile)
+        with tmpdir.as_cwd():
+            u.atoms.write(testfile)
 
-        u2 = mda.Universe(outfile)
+            u2 = mda.Universe(testfile)
 
-        assert_equal(u.atoms.positions,
-                     u2.atoms.positions)
+            assert_equal(u.atoms.positions,
+                         u2.atoms.positions)
 
     def test_roundtrip(self, u, outfile):
         # Write out a copy of the Universe, and compare this against the original

testsuite/MDAnalysisTests/coordinates/test_lammps.py

@@ -220,8 +220,13 @@ def test_molecule_tag(self, LAMMPSDATAWriter_molecule_tag):
                      err_msg="resids different after writing",)
 
 
-def test_datawriter_universe(tmpdir):
-    fn = str(tmpdir.join('out.data'))
+@pytest.mark.parametrize('filename', ['out.data', 'out.data.bz2'])
+def test_datawriter_universe(filename, tmpdir):
+    """
+    Test roundtrip on datawriter, and also checks compressed files
+    can be written (see #4159).
+    """
+    fn = str(tmpdir.join(filename))
 
     u = mda.Universe(LAMMPSdata_mini)
 

testsuite/MDAnalysisTests/coordinates/test_pdbqt.py

@@ -89,13 +89,17 @@ class TestPDBQTWriter(object):
     def outfile(self, tmpdir):
         return str(tmpdir) + 'out.pdbqt'
 
-    def test_roundtrip_writing_coords(self, outfile):
-        u = mda.Universe(PDBQT_input)
-        u.atoms.write(outfile)
-        u2 = mda.Universe(outfile)
+    @pytest.mark.parametrize('filename',
+        ['test.pdbqt', 'test.pdbqt.bz2', 'test.pdbqt.gz'])
+    def test_roundtrip_writing_coords(self, filename, tmpdir):
+
+        with tmpdir.as_cwd():
+            u = mda.Universe(PDBQT_input)
+            u.atoms.write(filename)
+            u2 = mda.Universe(filename)
 
-        assert_equal(u2.atoms.positions, u.atoms.positions,
-                     "Round trip does not preserve coordinates")
+            assert_equal(u2.atoms.positions, u.atoms.positions,
+                         "Round trip does not preserve coordinates")
 
     def test_roundtrip_formatting(self, outfile):
         # Compare formatting of first line

testsuite/MDAnalysisTests/coordinates/test_pqr.py

@@ -80,6 +80,15 @@ def universe():
 
     prec = 3
 
+    @pytest.mark.parametrize('filename',
+        ['test.pqr', 'test.pqr.bz2', 'test.pqr.gz'])
+    def test_simple_writer_roundtrip(self, universe, filename, tmpdir):
+        with tmpdir.as_cwd():
+            universe.atoms.write(filename)
+            u2 = mda.Universe(filename)
+            assert_equal(universe.atoms.positions,
+                         u2.atoms.positions)
+
     def test_writer_noChainID(self, universe, tmpdir):
         outfile = str(tmpdir.join('pqr-test.pqr'))