Serialize FileIO and TextIOWrapper and Universe

package/MDAnalysis/lib/util.py

@@ -231,7 +231,6 @@ class PathLike(object):
         pass
 
 
-
 def filename(name, ext=None, keep=False):
     """Return a new name that has suffix attached; replaces other extensions.
 
@@ -265,6 +264,98 @@ def filename(name, ext=None, keep=False):
     return name if isstream(name) else str(name)
 
 
+class FileIOPickable(io.FileIO):
+    """Stream for read a file
+
+    Picklable FiloIO class that only support read mode.
+
+    """
+    def __getstate__(self):
+        return self.tell(), self.name
+
+    def __setstate__(self, args):
+        name = args[1]
+        super().__init__(name)
+        self.seek(args[0])
+
+
+class BufferIOPickable(io.BufferedReader):
+    """A pickable buffer for a readable FilIO object
+
+    Wrap raw FileIOPickable inside
+
+    """
+    def __init__(self, raw):
+        super().__init__(raw)
+        self.raw_class = raw.__class__
+
+    def __getstate__(self):
+        return self.tell(), self.name, self.raw_class
+
+    def __setstate__(self, args):
+        name = args[1]
+        raw_class = args[2]
+        raw = raw_class(name)
+        super().__init__(raw)
+        self.seek(args[0])
+
+
+class TextIOPickable(io.TextIOWrapper):
+    """Character and line based layer over a pickable FileIO based object.
+
+    Example
+    -------
+        file = FileIOPickable('filename')
+        text_wrapped = TextIOPickable(file)
+    """
+    def __init__(self, raw):
+        super().__init__(raw)
+        self.raw_class = raw.__class__
+
+    def __getstate__(self):
+        return self.tell(), self.name, self.raw_class
+
+    def __setstate__(self, args):
+        name = args[1]
+        raw_class = args[2]
+        # raw_class is used for further expansion this functionality to 
+        # GZip files, which also requires a text wrapper. 
+        raw = raw_class(name)
+        super().__init__(raw)
+        self.seek(args[0])
+
+# not as comprehensive as built-in open func--no need for other args
+# only should be used for 'reading' modes
+def pickle_open(name, mode='rt'):
+    """Open file and return a stream with pickle function implemented.
+
+    Parameters
+    ----------
+    name : str;
+        a filename given a text or byte string.
+    mode: {'r', 'rt', 'rb'} (optional)
+        'r':  open for reading in text mode;
+        'rt': read in text mode (default);
+        'rb': read in binary mode;
+        raise ValueError with other modes.
+
+    Returns
+    -------
+    stream-like object
+
+    See Also
+    --------
+    :func:`anyopen`
+    """
+    if mode not in {'r', 'rt', 'rb'}:
+        raise ValueError("Only read mode ('r', 'rt', 'rb') files can be pickled.")
+    raw = FileIOPickable(name)
+    if mode == 'rb':
+        return BufferIOPickable(raw)
+    elif mode == 'rt' or mode == 'r':
+        return TextIOPickable(raw)
+
+
 @contextmanager
 def openany(datasource, mode='rt', reset=True):
     """Context manager for :func:`anyopen`.

testsuite/MDAnalysisTests/parallelism/test_multiprocessing.py

@@ -0,0 +1,53 @@
+# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
+# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8
+#
+# MDAnalysis --- https://www.mdanalysis.org
+# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
+# (see the file AUTHORS for the full list of names)
+#
+# Released under the GNU Public Licence, v2 or any higher version
+#
+# Please cite your use of MDAnalysis in published work:
+#
+# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
+# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
+# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
+# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
+# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
+# doi: 10.25080/majora-629e541a-00e
+#
+# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
+# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
+# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
+#
+
+from __future__ import absolute_import
+import sys
+import multiprocessing
+import numpy as np
+import pytest
+
+from MDAnalysis.lib.util import pickle_open
+from MDAnalysisTests.datafiles import (
+    PDB
+)
+
+from numpy.testing import assert_equal
+
+
+def textio_line(file, i):
+    return file.readlines()[i]
+
+
+@pytest.mark.xfail(sys.version_info < (3, 0), reason="pickle function not \
+                   working in python 2")
+def test_multiprocess_fileio():
+    p = multiprocessing.Pool(2)
+    PDB_file = pickle_open(PDB)
+    ref = PDB_file.readlines()[:4]
+    PDB_file.close()
+    PDB_file = pickle_open(PDB)
+    res = np.array([p.apply(textio_line, args=(PDB_file, i))
+                    for i in range(4)])
+    p.close()
+    assert_equal(res, ref)

testsuite/MDAnalysisTests/utils/test_pickleio.py

@@ -0,0 +1,58 @@
+# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
+# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8
+#
+# MDAnalysis --- https://www.mdanalysis.org
+# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
+# (see the file AUTHORS for the full list of names)
+#
+# Released under the GNU Public Licence, v2 or any higher version
+#
+# Please cite your use of MDAnalysis in published work:
+#
+# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
+# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
+# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
+# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
+# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
+# doi: 10.25080/majora-629e541a-00e
+#
+# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
+# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
+# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
+#
+from __future__ import absolute_import
+
+import pytest
+from numpy.testing import assert_equal
+
+import MDAnalysis.lib.util as util
+from MDAnalysis.tests.datafiles import PDB
+
+import pickle
+import sys
+
+
+@pytest.fixture(params=[
+    # filename mode
+    (PDB, 'r'),
+    (PDB, 'rt'),
+    (PDB, 'rb'),
+])
+def f(request):
+    filename, mode = request.param
+    return util.pickle_open(filename, mode)
+
+
+@pytest.mark.xfail(sys.version_info < (3, 0), reason="pickle function not \
+                     working in python 2")
+def test_iopickle(f):
+    f_pickled = pickle.loads(pickle.dumps(f))
+    assert_equal(f.readline(), f_pickled.readline())
+
+
+@pytest.mark.xfail(sys.version_info < (3, 0), reason="pickle function not \
+                     working in python 2")
+def test_offset(f):
+    f.readline()
+    f_pickled = pickle.loads(pickle.dumps(f))
+    assert_equal(f.tell(), f_pickled.tell())