Repository: joblib/joblib
Pull request #1681 (2025-02-27 14:44:34)

FIX joblib.load wrongly tries to load memmaps from a compressed dump

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/joblib/joblib/pull/1681?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 66.94%. Comparing base [(`73a78ac`)](https://app.codecov.io/gh/joblib/joblib/commit/73a78aca67afaa33aadbcc9795042dff1149a102?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) to head [(`4b06133`)](https://app.codecov.io/gh/joblib/joblib/commit/4b0613306994168f3f87f98f1339ea6bc9a3836f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib).

> :exclamation:  There is a different number of reports uploaded between BASE (73a78ac) and HEAD (4b06133). Click for more details.
> 
> <details><summary>HEAD has 23 uploads less than BASE</summary>
>
>| Flag | BASE (73a78ac) | HEAD (4b06133) |
>|------|------|------|
>||24|1|
></details>

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##             main    #1681       +/-   ##
===========================================
- Coverage   95.66%   66.94%   -28.73%     
===========================================
  Files          46       46               
  Lines        7724     7726        +2     
===========================================
- Hits         7389     5172     -2217     
- Misses        335     2554     +2219     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/joblib/joblib/pull/1681?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib).

lesteve: While you are at it maybe improve the naming of the function? `_read_fileobject` => `_get_fileobject_and_mmap` or something better?

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index 444e544..efe8504 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -33,6 +33,14 @@ In development
 - Drop support for PyPy.
   https://github.com/joblib/joblib/pull/1670
 
+- Fixed an issue affecting ``joblib.load`` calls with non-null ``mmap_mode``
+  parameter when loading compressed python objects. It wrongly attempted to load
+  with ``np.memmap`` anyway, resulting in python exceptions or corrupted data.
+  The result now properly use in-memory ``np.array`` arrays, in accordance with
+  the warnings that are emitted in this case.
+  https://github.com/joblib/joblib/pull/1681
+
+
 Release 1.4.2 -- 2024/05/02
 ---------------------------
 
diff --git a/joblib/numpy_pickle.py b/joblib/numpy_pickle.py
index 20a1bdc..8dc6b7a 100644
--- a/joblib/numpy_pickle.py
+++ b/joblib/numpy_pickle.py
@@ -38,7 +38,7 @@ from .numpy_pickle_utils import (
     Unpickler,
     _ensure_native_byte_order,
     _read_bytes,
-    _read_fileobject,
+    _validate_fileobject_and_memmap,
     _write_fileobject,
 )
 
@@ -678,16 +678,19 @@ def load(filename, mmap_mode=None):
     if hasattr(filename, "read"):
         fobj = filename
         filename = getattr(fobj, "name", "")
-        with _read_fileobject(fobj, filename, mmap_mode) as fobj:
+        with _validate_fileobject_and_memmap(fobj, filename, mmap_mode) as (fobj, _):
             obj = _unpickle(fobj)
     else:
         with open(filename, "rb") as f:
-            with _read_fileobject(f, filename, mmap_mode) as fobj:
+            with _validate_fileobject_and_memmap(f, filename, mmap_mode) as (
+                fobj,
+                validated_mmap_mode,
+            ):
                 if isinstance(fobj, str):
                     # 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)
+                obj = _unpickle(fobj, filename, validated_mmap_mode)
     return obj
diff --git a/joblib/numpy_pickle_utils.py b/joblib/numpy_pickle_utils.py
index f2034e0..c241715 100644
--- a/joblib/numpy_pickle_utils.py
+++ b/joblib/numpy_pickle_utils.py
@@ -128,7 +128,7 @@ def _buffered_write_file(fobj):
 
 
 @contextlib.contextmanager
-def _read_fileobject(fileobj, filename, mmap_mode=None):
+def _validate_fileobject_and_memmap(fileobj, filename, mmap_mode=None):
     """Utility function opening the right fileobject from a filename.
 
     The magic number is used to choose between the type of file object to open:
@@ -141,8 +141,6 @@ def _read_fileobject(fileobj, filename, mmap_mode=None):
     Parameters
     ----------
     fileobj: file object
-    compressor: str in {'zlib', 'gzip', 'bz2', 'lzma', 'xz', 'compat',
-                        'not-compressed'}
     filename: str
         filename path corresponding to the fileobj parameter.
     mmap_mode: str
@@ -152,11 +150,12 @@ def _read_fileobject(fileobj, filename, mmap_mode=None):
 
     Returns
     -------
-        a file like object
+        a tuple with a file like object, and the validated mmap_mode.
 
     """
     # Detect if the fileobj contains compressed data.
     compressor = _detect_compressor(fileobj)
+    validated_mmap_mode = mmap_mode
 
     if compressor == "compat":
         # Compatibility with old pickle mode: simply return the input
@@ -169,7 +168,7 @@ def _read_fileobject(fileobj, filename, mmap_mode=None):
             DeprecationWarning,
             stacklevel=2,
         )
-        yield filename
+        yield filename, validated_mmap_mode
     else:
         if compressor in _COMPRESSORS:
             # based on the compressor detected in the file, we open the
@@ -182,6 +181,7 @@ def _read_fileobject(fileobj, filename, mmap_mode=None):
         # mmap_mode cannot be used with compressed file or in memory buffers
         # such as io.BytesIO.
         if mmap_mode is not None:
+            validated_mmap_mode = None
             if isinstance(fileobj, io.BytesIO):
                 warnings.warn(
                     "In memory persistence is not compatible with "
@@ -202,8 +202,10 @@ def _read_fileobject(fileobj, filename, mmap_mode=None):
                     '"%(mmap_mode)s" flag will be ignored.' % locals(),
                     stacklevel=2,
                 )
+            else:
+                validated_mmap_mode = mmap_mode
 
-        yield fileobj
+        yield fileobj, validated_mmap_mode
 
 
 def _write_fileobject(filename, compress=("zlib", 3)):


Test patch:
diff --git a/joblib/test/test_numpy_pickle.py b/joblib/test/test_numpy_pickle.py
index 73ab91d..6c61710 100644
--- a/joblib/test/test_numpy_pickle.py
+++ b/joblib/test/test_numpy_pickle.py
@@ -279,13 +279,15 @@ def test_masked_array_persistence(tmpdir):
 def test_compress_mmap_mode_warning(tmpdir):
     # Test the warning in case of compress + mmap_mode
     rnd = np.random.RandomState(0)
-    a = rnd.random_sample(10)
+    obj = rnd.random_sample(10)
     this_filename = tmpdir.join("test.pkl").strpath
-    numpy_pickle.dump(a, this_filename, compress=1)
+    numpy_pickle.dump(obj, this_filename, compress=1)
     with warns(UserWarning) as warninfo:
-        numpy_pickle.load(this_filename, mmap_mode="r+")
+        reloaded_obj = numpy_pickle.load(this_filename, mmap_mode="r+")
     debug_msg = "\n".join([str(w) for w in warninfo])
     warninfo = [w.message for w in warninfo]
+    assert not isinstance(reloaded_obj, np.memmap)
+    np.testing.assert_array_equal(obj, reloaded_obj)
     assert len(warninfo) == 1, debug_msg
     assert (
         str(warninfo[0]) == 'mmap_mode "r+" is not compatible with compressed '