Repository: joblib/joblib

Issue:
In-memory caching
Sometimes reloading custom objects from the cache can be expensive (but still much faster than recomputing the whole object from scratch), in which case a two-level cache can be useful: results could be first kept in memory, and, after some timeout (LRU/etc.), pushed to the hard drive for longer-time persistence.
Is this something that joblib's architecture may easily allow?

Discussion:
aabadie: I think this is in a pretty good shape now, so changing the status to MRG.

A few comments though:
* the implementation is not perfect (e.g new `_COMPRESSOR_OBJS` list, etc) so I'd like to have someone else's opinion,
* forcing the decompressor to a wrong one (exemple: use gzip compression with dump but try to uncompress with bz2) will raise an exception. Should we handle this case ? My answer on this is 'no' as the message is quite meaningful but maybe other people won't agree with this,
* online documentation and CHANGES.rst still need to be updated but this can be done after a preliminary review.
aabadie: @lesteve, I addressed what we discussed IRL. Feel free to have a look when you have time.
aabadie: @lesteve, @ogrisel , any comment ?
aabadie: Just for the record, here is the kind of things now possible with this PR: https://gist.github.com/aabadie/074587354d97d872aff6abb65510f618
lesteve: Looks like you need to rebase on master.
lesteve: > Looks like you need to rebase on master.

Other than this LGTM.

aabadie: > Looks like you need to rebase on master.

Rebased, waiting for AppVeyor.
lesteve: Can you add an entry in CHANGES.rst? Maybe we want to update the docs and add a small section saying that we can use file objects.
aabadie: > Can you add an entry in CHANGES.rst? 

Sure

> Maybe we want to update the docs and add a small section saying that we can use file objects.

Good point. Something I forgot to add.
aabadie: @lesteve, I added some lines about persistence in file objects in the documentation and updated the CHANGES.rst file. Comments welcome :wink: 
lesteve: Can you rebase on master so we quickly double-check that flake8_diff.sh works fine on your PR?
aabadie: > Can you rebase on master so we quickly double-check that flake8_diff.sh works fine on your PR?

Travis is broken because of the new code in the doc. I'm on it.
aabadie: I had issues with doctest and the usage of file object in a context manager. I skipped the lines in the doc but that may not be a solution. Any better idea ?

Otherwise, I rebased on the lastest master (including the flake8_diff.sh recent changes) and it works just fine.
lesteve: > I had issues with doctest and the usage of file object in a context manager. I skipped the lines in the doc but that may not be a solution. Any better idea ?

What was the error? It would be great to understand where it comes from and fix it properly ...

> Otherwise, I rebased on the lastest master (including the flake8_diff.sh recent changes) and it works just fine.

Great, thanks a lot!
aabadie: > What was the error? It would be great to understand where it comes from and fix it properly ...

There were 2 things:
* `with open() as fo` returns the file object ([<io.blabla>]) : I fixed it using an ellipsis flags and by adding `[<...>]` after the `with` block
* in python 2.6, using gzip.GzipFile or bz2.BZ2File doesn't work if you don't use `contextlib.closing`. So I had to skip the file object compression example.
lesteve: > with open() as fo returns the file object ([]) : I fixed it using an ellipsis flags and by adding [<...>] after the with block

Hmmm I am not sure what we should do in this case:

```py
with gzip.GzipFile('/tmp/test.gz', 'wb') as f:             
    res = joblib.dump({'a': [1, 2, 3], 'b': 'asadf'}, f)
```

Should `res` be:
* a single element list containing containing `f`
* a single element list containing containing f.name to be more consistent with when dump is called with a filename
* because this is a new feature, can we do what we want, in particular returning a list may not make sense, since the list was there for historical reasons (main pickle + companion files for each numpy array)

> in python 2.6, using gzip.GzipFile or bz2.BZ2File doesn't work if you don't use contextlib.closing. So I had to skip the file object compression example.

Pfff :sob: is there a way we can just skip doctests in Python 2.6?
lesteve: > Pfff :sob: is there a way we can just skip doctests in Python 2.6?

Seems like scikit-learn has some "*fixture*" files in doc to skip the doctests. You can probably have a look there. I believe there is something in setup.cfg as well to work with nose and doctest-fixtures.
aabadie: > Should res be:
>  *  a single element list containing containing f
>  *  a single element list containing containing f.name to be more consistent with when dump is called > with a filename
>  *  because this is a new feature, can we do what we want, in particular returning a list may not make > sense, since the list was there for historical reasons (main pickle + companion files for each numpy array)

I think returning nothing in case a file object is given as input is reasonable because this means the user knows in advance the destination of the dump and the returned list is here, as you said, for historical reasons.

aabadie: > Seems like scikit-learn has some *fixture* files in the doc folder to skip the doctests

Nice, it works like a charm. The `setup.cfg` file in joblib was already correctly configured.

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index 082f60c..eed5f3d 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -4,6 +4,11 @@ Latest changes
 Release 0.10.0
 --------------
 
+Alexandre Abadie
+
+    ENH: joblib.dump/load now accept file-like objects besides filenames.
+    https://github.com/joblib/joblib/pull/351 for more details.
+
 Niels Zeilemaker and Olivier Grisel
 
     Refactored joblib.Parallel to enable the registration of custom
diff --git a/benchmarks/bench_auto_batching.py b/benchmarks/bench_auto_batching.py
index 6046099..d94bc95 100644
--- a/benchmarks/bench_auto_batching.py
+++ b/benchmarks/bench_auto_batching.py
@@ -53,11 +53,10 @@ def bench_short_tasks(task_times, n_jobs=2, batch_size="auto",
         p(delayed(sleep_noop)(max(t, 0), input_data, output_data_size)
           for t in task_times)
         duration = time.time() - t0
-        effective_batch_size = getattr(p, '_effective_batch_size',
+        effective_batch_size = getattr(p._backend, '_effective_batch_size',
                                        p.batch_size)
-
-    print('Completed %d tasks in %0.3fs, final batch_size=%d\n'
-          % (len(task_times), duration, effective_batch_size))
+    print('Completed {} tasks in {:3f}s, final batch_size={}\n'.format(
+        len(task_times), duration, effective_batch_size))
     return duration, effective_batch_size
 
 
diff --git a/doc/persistence.rst b/doc/persistence.rst
index d261795..f040428 100644
--- a/doc/persistence.rst
+++ b/doc/persistence.rst
@@ -46,6 +46,18 @@ We can then load the object from the file::
   [('a', [1, 2, 3]), ('b', array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))]
 
 
+Persistence in file objects
+===========================
+
+Instead of filenames, `dump` and `load` functions also accept file objects:
+
+  >>> with open(filename, 'wb') as fo:  # doctest: +ELLIPSIS
+  ...    joblib.dump(to_persist, fo)
+  >>> with open(filename, 'rb') as fo:  # doctest: +ELLIPSIS
+  ...    joblib.load(fo)
+  [('a', [1, 2, 3]), ('b', array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))]
+
+
 Compressed joblib pickles
 =========================
 
@@ -79,6 +91,15 @@ are 'gzip', 'bz2', 'lzma' and 'xz':
 
     Lzma and Xz compression methods are only available for python versions >= 3.3.
 
+Compressor files provided by the python standard library can also be used to
+compress pickle, e.g ``gzip.GzipFile``, ``bz2.BZ2File``, ``lzma.LZMAFile``:
+    >>> # Dumping in a gzip.GzipFile object using a compression level of 3.
+    >>> import gzip
+    >>> with gzip.GzipFile(filename + '.gz', 'wb', compresslevel=3) as fo:  # doctest: +ELLIPSIS
+    ...    joblib.dump(to_persist, fo)
+    >>> with gzip.GzipFile(filename + '.gz', 'rb') as fo:  # doctest: +ELLIPSIS
+    ...    joblib.load(fo)
+    [('a', [1, 2, 3]), ('b', array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))]
 
 More details can be found in the :func:`joblib.dump` and
 :func:`joblib.load` documentation.
diff --git a/doc/persistence_fixture.py b/doc/persistence_fixture.py
new file mode 100644
index 0000000..503d0bc
--- /dev/null
+++ b/doc/persistence_fixture.py
@@ -0,0 +1,13 @@
+"""Fixture module to skip the persistence doctest with python 2.6."""
+
+from nose import SkipTest
+from joblib import _compat
+
+
+def setup_module(module):
+    """Setup module."""
+    if _compat.PY26:
+        # gzip.GZipFile and bz2.BZ2File compressor classes cannot be used
+        # within a context manager (e.g in a `with` block) in python 2.6 so
+        # we skip doctesting of persistence documentation in this case.
+        raise SkipTest("Skipping persistence doctest in Python 2.6")
diff --git a/joblib/numpy_pickle.py b/joblib/numpy_pickle.py
index f029582..0cf88a2 100644
--- a/joblib/numpy_pickle.py
+++ b/joblib/numpy_pickle.py
@@ -397,6 +397,9 @@ def dump(value, filename, compress=0, protocol=None, cache_size=None):
     if Path is not None and isinstance(filename, Path):
         filename = str(filename)
 
+    is_filename = isinstance(filename, _basestring)
+    is_fileobj = hasattr(filename, "write")
+
     compress_method = 'zlib'  # zlib is the default compression method.
     if compress is True:
         # By default, if compress is enabled, we want to be using 3 by default
@@ -424,14 +427,16 @@ def dump(value, filename, compress=0, protocol=None, cache_size=None):
             'Non valid compression method given: "{0}". Possible values are '
             '{1}.'.format(compress_method, _COMPRESSORS))
 
-    if not isinstance(filename, _basestring):
+    if not is_filename and not is_fileobj:
         # People keep inverting arguments, and the resulting error is
         # incomprehensible
         raise ValueError(
-            'Second argument should be a filename, %s (type %s) was given'
+            'Second argument should be a filename or a file-like object, '
+            '%s (type %s) was given.'
             % (filename, type(filename))
         )
-    elif not isinstance(compress, tuple):
+
+    if is_filename and not isinstance(compress, tuple):
         # In case no explicit compression was requested using both compression
         # method and level in a tuple and the filename has an explicit
         # extension, we select the corresponding compressor.
@@ -473,14 +478,54 @@ def dump(value, filename, compress=0, protocol=None, cache_size=None):
         with _write_fileobject(filename, compress=(compress_method,
                                                    compress_level)) as f:
             NumpyPickler(f, protocol=protocol).dump(value)
-
-    else:
+    elif is_filename:
         with open(filename, 'wb') as f:
             NumpyPickler(f, protocol=protocol).dump(value)
+    else:
+        NumpyPickler(filename, protocol=protocol).dump(value)
+
+    # If the target container is a file object, nothing is returned.
+    if is_fileobj:
+        return
 
+    # For compatibility, the list of created filenames (e.g with one element
+    # after 0.10.0) is returned by default.
     return [filename]
 
 
+def _unpickle(fobj, filename="", mmap_mode=None):
+    """Internal unpickling function."""
+    # We are careful to open the file handle early and keep it open to
+    # avoid race-conditions on renames.
+    # That said, if data is stored in companion files, which can be
+    # the case with the old persistence format, moving the directory
+    # will create a race when joblib tries to access the companion
+    # files.
+    unpickler = NumpyUnpickler(filename, fobj, mmap_mode=mmap_mode)
+    obj = None
+    try:
+        obj = unpickler.load()
+        if unpickler.compat_mode:
+            warnings.warn("The file '%s' has been generated with a "
+                          "joblib version less than 0.10. "
+                          "Please regenerate this pickle file."
+                          % filename,
+                          DeprecationWarning, stacklevel=3)
+    except UnicodeDecodeError as exc:
+        # More user-friendly error message
+        if PY3_OR_LATER:
+            new_exc = ValueError(
+                'You may be trying to read with '
+                'python 3 a joblib pickle generated with python 2. '
+                'This feature is not supported by joblib.')
+            new_exc.__cause__ = exc
+            raise new_exc
+        # Reraise exception with Python 2
+        raise
+
+    return obj
+
+
 def load(filename, mmap_mode=None):
     """Reconstruct a Python object from a file persisted with joblib.dump.
 
@@ -514,37 +559,19 @@ def load(filename, mmap_mode=None):
     """
     if Path is not None and isinstance(filename, Path):
         filename = str(filename)
-    with open(filename, 'rb') as f:
-        with _read_fileobject(f, filename, mmap_mode) as f:
-            if isinstance(f, _basestring):
-                # if the returned file object is a string, this means we try
-                # to load a pickle file generated with an version of Joblib
-                # so we load it with joblib compatibility function.
-                return load_compatibility(f)
-
-            # We are careful to open the file handle early and keep it open to
-            # avoid race-conditions on renames.
-            # That said, if data is stored in companion files, which can be
-            # the case with the old persistence format, moving the directory
-            # will create a race when joblib tries to access the companion
-            # files.
-            unpickler = NumpyUnpickler(filename, f, mmap_mode=mmap_mode)
-            try:
-                obj = unpickler.load()
-                if unpickler.compat_mode:
-                    warnings.warn("The file '%s' has been generated with a "
-                                  "joblib version less than 0.10. "
-                                  "Please regenerate this pickle file."
-                                  % filename,
-                                  DeprecationWarning, stacklevel=2)
-            except UnicodeDecodeError as exc:
-                # More user-friendly error message
-                if PY3_OR_LATER:
-                    new_exc = ValueError(
-                        'You may be trying to read with '
-                        'python 3 a joblib pickle generated with python 2. '
-                        'This feature is not supported by joblib.')
-                    new_exc.__cause__ = exc
-                    raise new_exc
+
+    if hasattr(filename, "read") and hasattr(filename, "seek"):
+        with _read_fileobject(filename, "", mmap_mode) as fobj:
+            obj = _unpickle(fobj)
+    else:
+        with open(filename, 'rb') as f:
+            with _read_fileobject(f, filename, mmap_mode) as fobj:
+                if isinstance(fobj, _basestring):
+                    # if the returned file object is a string, this means we
+                    # try to load a pickle file generated with an version of
+                    # Joblib so we load it with joblib compatibility function.
+                    return load_compatibility(fobj)
+
+                obj = _unpickle(fobj, filename, mmap_mode)
 
     return obj
diff --git a/joblib/numpy_pickle_utils.py b/joblib/numpy_pickle_utils.py
index 1755b28..ee879a6 100644
--- a/joblib/numpy_pickle_utils.py
+++ b/joblib/numpy_pickle_utils.py
@@ -8,6 +8,7 @@ import pickle
 import sys
 import io
 import zlib
+import gzip
 import bz2
 import warnings
 import contextlib
@@ -49,6 +50,9 @@ _LZMA_PREFIX = b'\x5d\x00'
 
 # Supported compressors
 _COMPRESSORS = ('zlib', 'bz2', 'lzma', 'xz', 'gzip')
+_COMPRESSOR_CLASSES = [gzip.GzipFile, bz2.BZ2File]
+if lzma is not None:
+    _COMPRESSOR_CLASSES.append(lzma.LZMAFile)
 
 # The max magic number length of supported compression file types.
 _MAX_PREFIX_LEN = max(len(prefix)
@@ -147,20 +151,42 @@ def _read_fileobject(fileobj, filename, mmap_mode=None):
     """
     # Detect if the fileobj contains compressed data.
     compressor = _detect_compressor(fileobj)
+    if isinstance(fileobj, tuple(_COMPRESSOR_CLASSES)):
+        compressor = fileobj.__class__.__name__
     if compressor == 'compat':
+        # Compatibility with old pickle mode: simply return the input
+        # filename "as-is" and let the compatibility function be called by the
+        # caller.
         warnings.warn("The file '%s' has been generated with a joblib "
                       "version less than 0.10. "
                       "Please regenerate this pickle file." % filename,
                       DeprecationWarning, stacklevel=2)
         yield filename
     else:
-        if compressor in _COMPRESSORS and mmap_mode is not None:
+        # Checking if incompatible load parameters with the type of file:
+        # mmap_mode cannot be used with compressed file or in memory buffers
+        # such as io.BytesIO.
+        if ((compressor in _COMPRESSORS or
+                isinstance(fileobj, tuple(_COMPRESSOR_CLASSES))) and
+                mmap_mode is not None):
             warnings.warn('File "%(filename)s" is compressed using '
                           '"%(compressor)s" which is not compatible with '
-                          'mmap_mode "%(mmap_mode)s" flag passed.'
-                          % locals(), DeprecationWarning, stacklevel=2)
-
-        if compressor == 'zlib':
+                          'mmap_mode "%(mmap_mode)s" flag passed. mmap_mode '
+                          'option will be ignored.'
+                          % locals(), stacklevel=2)
+        if isinstance(fileobj, io.BytesIO) and mmap_mode is not None:
+            warnings.warn('In memory persistence is not compatible with '
+                          'mmap_mode "%(mmap_mode)s" flag passed. mmap_mode '
+                          'option will be ignored.'
+                          % locals(), stacklevel=2)
+
+        # if the passed fileobj is in the supported list of decompressor
+        # objects (GzipFile, BZ2File, LzmaFile), we simply return it.
+        if isinstance(fileobj, tuple(_COMPRESSOR_CLASSES)):
+            yield fileobj
+        # otherwise, based on the compressor detected in the file, we open the
+        # correct decompressor file object, wrapped in a buffer.
+        elif compressor == 'zlib':
             yield _buffered_read_file(BinaryZlibFile(fileobj, 'rb'))
         elif compressor == 'gzip':
             yield _buffered_read_file(BinaryGzipFile(fileobj, 'rb'))
@@ -180,6 +206,7 @@ def _read_fileobject(fileobj, filename, mmap_mode=None):
                                           "python ({0}.{1})"
                                           .format(sys.version_info[0],
                                                   sys.version_info[1]))
+        # No compression detected => returning the input file object (open)
         else:
             yield fileobj
 


Test patch:
diff --git a/joblib/test/test_numpy_pickle.py b/joblib/test/test_numpy_pickle.py
index 19a5e95..f3708fb 100644
--- a/joblib/test/test_numpy_pickle.py
+++ b/joblib/test/test_numpy_pickle.py
@@ -13,6 +13,7 @@ import warnings
 import nose
 import gzip
 import zlib
+import bz2
 import pickle
 from contextlib import closing
 
@@ -25,7 +26,7 @@ from joblib.testing import assert_raises_regex
 from joblib import numpy_pickle
 from joblib.test import data
 
-from joblib._compat import PY3_OR_LATER
+from joblib._compat import PY3_OR_LATER, PY26
 from joblib.numpy_pickle_utils import _IO_BUFFER_SIZE, BinaryZlibFile
 from joblib.numpy_pickle_utils import _detect_compressor, _COMPRESSORS
 
@@ -334,14 +335,15 @@ def test_compress_mmap_mode_warning():
         numpy_pickle.load(this_filename, mmap_mode='r+')
         nose.tools.assert_equal(len(caught_warnings), 1)
         for warn in caught_warnings:
-            nose.tools.assert_equal(warn.category, DeprecationWarning)
+            nose.tools.assert_equal(warn.category, UserWarning)
             nose.tools.assert_equal(warn.message.args[0],
                                     'File "%(filename)s" is compressed using '
                                     '"%(compressor)s" which is not compatible '
                                     'with mmap_mode "%(mmap_mode)s" flag '
-                                    'passed.' % {'filename': this_filename,
-                                                 'mmap_mode': 'r+',
-                                                 'compressor': 'zlib'})
+                                    'passed. mmap_mode option will be '
+                                    'ignored.' % {'filename': this_filename,
+                                                  'mmap_mode': 'r+',
+                                                  'compressor': 'zlib'})
 
 
 @with_numpy
@@ -681,6 +683,120 @@ def test_compression_using_file_extension():
             os.remove(dump_fname)
 
 
+@with_numpy
+def test_file_handle_persistence():
+    objs = [np.random.random((10, 10)),
+            "some data",
+            np.matrix([0, 1, 2])]
+    fobjs = [open]
+    if not PY26:
+        fobjs += [bz2.BZ2File, gzip.GzipFile]
+    if PY3_OR_LATER:
+        import lzma
+        fobjs += [lzma.LZMAFile]
+    filename = env['filename'] + str(random.randint(0, 1000))
+
+    for obj in objs:
+        for fobj in fobjs:
+            with fobj(filename, 'wb') as f:
+                numpy_pickle.dump(obj, f)
+
+            # using the same decompressor prevents from internally
+            # decompress again.
+            with fobj(filename, 'rb') as f:
+                obj_reloaded = numpy_pickle.load(f)
+
+            # when needed, the correct decompressor should be used when
+            # passing a raw file handle.
+            with open(filename, 'rb') as f:
+                obj_reloaded_2 = numpy_pickle.load(f)
+
+            if isinstance(obj, np.ndarray):
+                np.testing.assert_array_equal(obj_reloaded, obj)
+                np.testing.assert_array_equal(obj_reloaded_2, obj)
+            else:
+                nose.tools.assert_equal(obj_reloaded, obj)
+                nose.tools.assert_equal(obj_reloaded_2, obj)
+
+            os.remove(filename)
+
+
+@with_numpy
+def test_in_memory_persistence():
+    objs = [np.random.random((10, 10)),
+            "some data",
+            np.matrix([0, 1, 2])]
+    for obj in objs:
+        f = io.BytesIO()
+        numpy_pickle.dump(obj, f)
+        obj_reloaded = numpy_pickle.load(f)
+        if isinstance(obj, np.ndarray):
+            np.testing.assert_array_equal(obj_reloaded, obj)
+        else:
+            nose.tools.assert_equal(obj_reloaded, obj)
+
+
+@with_numpy
+def test_file_handle_persistence_mmap():
+    obj = np.random.random((10, 10))
+    filename = env['filename'] + str(random.randint(0, 1000))
+
+    with open(filename, 'wb') as f:
+        numpy_pickle.dump(obj, f)
+
+    with open(filename, 'rb') as f:
+        obj_reloaded = numpy_pickle.load(f, mmap_mode='r+')
+
+    np.testing.assert_array_equal(obj_reloaded, obj)
+
+
+@with_numpy
+def test_file_handle_persistence_compressed_mmap():
+    obj = np.random.random((10, 10))
+    filename = env['filename'] + str(random.randint(0, 1000))
+
+    with open(filename, 'wb') as f:
+        numpy_pickle.dump(obj, f, compress=('gzip', 3))
+
+    with closing(gzip.GzipFile(filename, 'rb')) as f:
+        with warnings.catch_warnings(record=True) as caught_warnings:
+            warnings.simplefilter("always")
+            numpy_pickle.load(f, mmap_mode='r+')
+            nose.tools.assert_equal(len(caught_warnings), 1)
+            for warn in caught_warnings:
+                nose.tools.assert_equal(warn.category, UserWarning)
+                nose.tools.assert_equal(warn.message.args[0],
+                                        'File "%(filename)s" is compressed '
+                                        'using "%(compressor)s" which is not '
+                                        'compatible with mmap_mode '
+                                        '"%(mmap_mode)s" flag '
+                                        'passed. mmap_mode option will be '
+                                        'ignored.' %
+                                        {'filename': "",
+                                         'mmap_mode': 'r+',
+                                         'compressor': 'GzipFile'})
+
+
+@with_numpy
+def test_file_handle_persistence_in_memory_mmap():
+    obj = np.random.random((10, 10))
+    buf = io.BytesIO()
+
+    numpy_pickle.dump(obj, buf)
+
+    with warnings.catch_warnings(record=True) as caught_warnings:
+        warnings.simplefilter("always")
+        numpy_pickle.load(buf, mmap_mode='r+')
+        nose.tools.assert_equal(len(caught_warnings), 1)
+        for warn in caught_warnings:
+            nose.tools.assert_equal(warn.category, UserWarning)
+            nose.tools.assert_equal(warn.message.args[0],
+                                    'In memory persistence is not compatible '
+                                    'with mmap_mode "%(mmap_mode)s" '
+                                    'flag passed. mmap_mode option will be '
+                                    'ignored.' % {'mmap_mode': 'r+'})
+
+
 def test_binary_zlibfile():
     filename = env['filename'] + str(random.randint(0, 1000))