Serialize FileIO and TextIOWrapper and Universe

package/CHANGELOG

@@ -13,9 +13,11 @@ The rules for this file:
   * release numbers follow "Semantic Versioning" http://semver.org
 
 ------------------------------------------------------------------------------
+
 ??/??/?? tylerjereddy, richardjgowers, IAlibay, hmacdope, orbeckst, cbouy,
          lilyminium, daveminh, jbarnoud, yuxuanzhuang
 
+
   * 2.0.0
 
 Fixes
@@ -40,6 +42,9 @@ Enhancements
   * Added computation of Mean Squared Displacements (#2438, PR #2619)
   * Improved performances when parsing TPR files (PR #2804)
   * Added converter between Cartesian and Bond-Angle-Torsion coordinates (PR #2668)
+  * Added lib.pickle_file_io module for pickling file handlers. (PR #2723)
+  * Added pickle function to `Universe` and all Readers (without transformation)
+    (PR #2723)
 
 Changes
   * deprecated NumPy type aliases have been replaced with their actual types
@@ -59,6 +64,7 @@ Changes
   * Removes deprecated ProgressMeter (Issue #2739)
   * Removes deprecated MDAnalysis.units.N_Avogadro (PR #2737)
   * Dropped Python 2 support
+  * Set Python 3.6 as the minimum supported version (Issue #2541)
 
 Deprecations
 

package/MDAnalysis/coordinates/DLPoly.py

@@ -32,6 +32,7 @@
 
 from . import base
 from . import core
+from ..lib import util
 
 _DLPOLY_UNITS = {'length': 'Angstrom', 'velocity': 'Angstrom/ps', 'time': 'ps'}
 
@@ -149,7 +150,7 @@ def __init__(self, filename, **kwargs):
         super(HistoryReader, self).__init__(filename, **kwargs)
 
         # "private" file handle
-        self._file = open(self.filename, 'r')
+        self._file = util.anyopen(self.filename, 'r')
         self.title = self._file.readline().strip()
         self._levcfg, self._imcon, self.n_atoms = np.int64(self._file.readline().split()[:3])
         self._has_vels = True if self._levcfg > 0 else False

package/MDAnalysis/coordinates/GSD.py

@@ -44,13 +44,20 @@
 .. autoclass:: GSDReader
    :inherited-members:
 
+.. autoclass:: GSDPicklable
+   :members:
+
+.. autofunction:: gsd_pickle_open
+
 """
 import numpy as np
-import os
+import gsd
+import gsd.fl as fl
 import gsd.hoomd
 
 from . import base
 
+
 class GSDReader(base.ReaderBase):
     """Reader for the GSD format.
 
@@ -69,6 +76,10 @@ def __init__(self, filename, **kwargs):
 
 
         .. versionadded:: 0.17.0
+        .. versionchanged:: 2.0.0
+            Now use a picklable :class:`gsd.hoomd.HOOMDTrajectory`--
+            :class:`GSDPicklable`
+
         """
         super(GSDReader, self).__init__(filename, **kwargs)
         self.filename = filename
@@ -77,10 +88,10 @@ def __init__(self, filename, **kwargs):
         self.ts = self._Timestep(self.n_atoms, **self._ts_kwargs)
         self._read_next_timestep()
 
-    def open_trajectory(self) :
+    def open_trajectory(self):
         """opens the trajectory file using gsd.hoomd module"""
         self._frame = -1
-        self._file = gsd.hoomd.open(self.filename,mode='rb')
+        self._file = gsd_pickle_open(self.filename, mode='rb')
 
     def close(self):
         """close reader"""
@@ -97,7 +108,7 @@ def _reopen(self):
         self.open_trajectory()
 
     def _read_frame(self, frame):
-        try :
+        try:
             myframe = self._file[frame]
         except IndexError:
             raise IOError from None
@@ -111,20 +122,146 @@ def _read_frame(self, frame):
 
         # set frame box dimensions
         self.ts.dimensions = myframe.configuration.box
-        for i in range(3,6) :
+        for i in range(3, 6):
             self.ts.dimensions[i] = np.arccos(self.ts.dimensions[i]) * 180.0 / np.pi
 
         # set particle positions
         frame_positions = myframe.particles.position
         n_atoms_now = frame_positions.shape[0]
-        if n_atoms_now != self.n_atoms :
+        if n_atoms_now != self.n_atoms:
             raise ValueError("Frame %d has %d atoms but the initial frame has %d"
                 " atoms. MDAnalysis in unable to deal with variable"
                 " topology!"%(frame, n_atoms_now, self.n_atoms))
-        else :
+        else:
             self.ts.positions = frame_positions
         return self.ts
 
-    def _read_next_timestep(self) :
+    def _read_next_timestep(self):
         """read next frame in trajectory"""
         return self._read_frame(self._frame + 1)
+
+
+class GSDPicklable(gsd.hoomd.HOOMDTrajectory):
+    """Hoomd GSD file object (read-only) that can be pickled.
+
+    This class provides a file-like object (as by :func:`gsd.hoomd.open`,
+    namely :class:`gsd.hoodm.HOOMDTrajectory`) that, unlike file objects,
+    can be pickled. Only read mode is supported.
+
+    When the file is pickled, filename and mode of :class:`gsd.fl.GSDFile` in
+    the file are saved. On unpickling, the file is opened by filename.
+    This means that for a successful unpickle, the original file still has to
+    be accessible with its filename.
+
+    Note
+    ----
+    Open hoomd GSD files with `gsd_pickle_open`.
+    After pickling, the current frame is reset. `universe.trajectory[i]` has
+    to be used to return to its original frame.
+
+    Parameters
+    ----------
+    file: :class:`gsd.fl.GSDFile`
+        File to access.
+
+    Example
+    -------
+    ::
+
+        gsdfileobj = gsd.fl.open(name=filename,
+                                     mode='rb',
+                                     application='gsd.hoomd '+gsd.__version__,
+                                     schema='hoomd',
+                                     schema_version=[1, 3]
+        file = GSDPicklable(gsdfileobj)
+        file_pickled = pickle.loads(pickle.dumps(file))
+
+    See Also
+    ---------
+    :func:`MDAnalysis.lib.picklable_file_io.FileIOPicklable`
+    :func:`MDAnalysis.lib.picklable_file_io.BufferIOPicklable`
+    :func:`MDAnalysis.lib.picklable_file_io.TextIOPicklable`
+    :func:`MDAnalysis.lib.picklable_file_io.GzipPicklable`
+    :func:`MDAnalysis.lib.picklable_file_io.BZ2Picklable`
+
+
+    .. versionadded:: 2.0.0
+    """
+    def __getstate__(self):
+        return self.file.name, self.file.mode
+
+    def __setstate__(self, args):
+        gsd_version = gsd.__version__
+        schema_version = [1, 4] if gsd_version >= '1.9.0' else [1, 3]
+        gsdfileobj = fl.open(name=args[0],
+                             mode=args[1],
+                             application='gsd.hoomd ' + gsd_version,
+                             schema='hoomd',
+                             schema_version=schema_version)
+        self.__init__(gsdfileobj)
+
+
+def gsd_pickle_open(name, mode='rb'):
+    """Open hoomd schema GSD file with pickle function implemented.
+
+    This function returns either GSDPicklable object. It can be used as a
+    context manager, and replace the built-in :func:`gsd.hoomd.open` function
+    in read mode that only returns an unpicklable file object.
+
+    Schema version will depend on the version of gsd module.
+
+    Note
+    ----
+    Can be only used with read mode.
+
+    Parameters
+    ----------
+    name : str
+        a filename given a text or byte string.
+    mode: {'r', 'rb'} (optional)
+        'r', 'rb':  open for reading;
+
+    Returns
+    -------
+    stream-like object: GSDPicklable
+
+    Raises
+    ------
+    ValueError
+        if `mode` is not one of the allowed read modes
+
+    Examples
+    -------
+    open as context manager::
+
+        with gsd_pickle_open('filename') as f:
+            line = f.readline()
+
+    open as function::
+
+        f = gsd_pickle_open('filename')
+        line = f.readline()
+        f.close()
+
+    See Also
+    --------
+    :func:`MDAnalysis.lib.util.anyopen`
+    :func:`MDAnalysis.lib.picklable_file_io.pickle_open`
+    :func:`MDAnalysis.lib.picklable_file_io.bz2_pickle_open`
+    :func:`MDAnalysis.lib.picklable_file_io.gzip_pickle_open`
+    :func:`gsd.hoomd.open`
+
+
+    .. versionadded:: 2.0.0
+    """
+    gsd_version = gsd.__version__
+    schema_version = [1, 4] if gsd_version >= '1.9.0' else [1, 3]
+    if mode not in {'r', 'rb'}:
+        raise ValueError("Only read mode ('r', 'rb') \
+                         files can be pickled.")
+    gsdfileobj = fl.open(name=name,
+                         mode=mode,
+                         application='gsd.hoomd ' + gsd_version,
+                         schema='hoomd',
+                         schema_version=schema_version)
+    return GSDPicklable(gsdfileobj)

package/MDAnalysis/coordinates/TRJ.py

@@ -85,6 +85,8 @@
 .. autoclass:: NCDFWriter
    :members:
 
+.. autoclass:: NCDFPicklable
+   :members:
 
 .. _ascii-trajectories:
 
@@ -158,7 +160,6 @@
 import MDAnalysis
 from . import base
 from ..lib import util
-
 logger = logging.getLogger("MDAnalysis.coordinates.AMBER")
 
 
@@ -450,6 +451,9 @@ class NCDFReader(base.ReaderBase):
     .. versionchanged:: 1.0.0
        Support for reading `degrees` units for `cell_angles` has now been
        removed (Issue #2327)
+    .. versionchanged:: 2.0.0
+       Now use a picklable :class:`scipy.io.netcdf.netcdf_file`--
+       :class:`NCDFPicklable`.
 
     """
 
@@ -469,8 +473,8 @@ def __init__(self, filename, n_atoms=None, mmap=None, **kwargs):
 
         super(NCDFReader, self).__init__(filename, **kwargs)
 
-        self.trjfile = scipy.io.netcdf.netcdf_file(self.filename,
-                                                   mmap=self._mmap)
+        self.trjfile = NCDFPicklable(self.filename,
+                                     mmap=self._mmap)
 
         # AMBER NetCDF files should always have a convention
         try:
@@ -1075,3 +1079,56 @@ def close(self):
         if self.trjfile is not None:
             self.trjfile.close()
             self.trjfile = None
+
+
+class NCDFPicklable(scipy.io.netcdf.netcdf_file):
+    """NetCDF file object (read-only) that can be pickled.
+
+    This class provides a file-like object (as returned by
+    :class:`scipy.io.netcdf.netcdf_file`) that,
+    unlike standard Python file objects,
+    can be pickled. Only read mode is supported.
+
+    When the file is pickled, filename and mmap of the open file handle in
+    the file are saved. On unpickling, the file is opened by filename,
+    and the mmap file is loaded.
+    This means that for a successful unpickle, the original file still has to
+    be accessible with its filename.
+
+    Parameters
+    ----------
+    filename : str or file-like
+        a filename given a text or byte string.
+    mmap : None or bool, optional
+        Whether to mmap `filename` when reading. True when `filename`
+        is a file name, False when `filename` is a file-like object.
+
+    Example
+    -------
+    ::
+
+        f = NCDFPicklable(NCDF)
+        print(f.variables['coordinates'].data)
+        f.close()
+
+    can also be used as context manager::
+
+        with NCDFPicklable(NCDF) as f:
+            print(f.variables['coordinates'].data)
+
+    See Also
+    ---------
+    :class:`MDAnalysis.lib.picklable_file_io.FileIOPicklable`
+    :class:`MDAnalysis.lib.picklable_file_io.BufferIOPicklable`
+    :class:`MDAnalysis.lib.picklable_file_io.TextIOPicklable`
+    :class:`MDAnalysis.lib.picklable_file_io.GzipPicklable`
+    :class:`MDAnalysis.lib.picklable_file_io.BZ2Picklable`
+
+
+    .. versionadded:: 2.0.0
+    """
+    def __getstate__(self):
+        return self.filename, self.use_mmap
+
+    def __setstate__(self, args):
+        self.__init__(args[0], mmap=args[1])