Merge pull request numpy#27335 from stefanv/savez-allow-pickle

ENH: Add ``allow_pickle`` flag to ``savez``

Author: charris

numpy/lib/_npyio_impl.py

@@ -588,13 +588,13 @@ def save(file, arr, allow_pickle=True, fix_imports=np._NoValue):
                            pickle_kwargs=dict(fix_imports=fix_imports))
 
 
-def _savez_dispatcher(file, *args, **kwds):
+def _savez_dispatcher(file, *args, allow_pickle=True, **kwds):
     yield from args
     yield from kwds.values()
 
 
 @array_function_dispatch(_savez_dispatcher)
-def savez(file, *args, **kwds):
+def savez(file, *args, allow_pickle=True, **kwds):
     """Save several arrays into a single file in uncompressed ``.npz`` format.
 
     Provide arrays as keyword arguments to store them under the
@@ -614,6 +614,14 @@ def savez(file, *args, **kwds):
         Arrays to save to the file. Please use keyword arguments (see
         `kwds` below) to assign names to arrays.  Arrays specified as
         args will be named "arr_0", "arr_1", and so on.
+    allow_pickle : bool, optional
+        Allow saving object arrays using Python pickles. Reasons for
+        disallowing pickles include security (loading pickled data can execute
+        arbitrary code) and portability (pickled objects may not be loadable
+        on different Python installations, for example if the stored objects
+        require libraries that are not available, and not all pickled data is
+        compatible between different versions of Python).
+        Default: True
     kwds : Keyword arguments, optional
         Arrays to save to the file. Each array will be saved to the
         output file with its corresponding keyword name.
@@ -678,16 +686,16 @@ def savez(file, *args, **kwds):
     array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
 
     """
-    _savez(file, args, kwds, False)
+    _savez(file, args, kwds, False, allow_pickle=allow_pickle)
 
 
-def _savez_compressed_dispatcher(file, *args, **kwds):
+def _savez_compressed_dispatcher(file, *args, allow_pickle=True, **kwds):
     yield from args
     yield from kwds.values()
 
 
 @array_function_dispatch(_savez_compressed_dispatcher)
-def savez_compressed(file, *args, **kwds):
+def savez_compressed(file, *args, allow_pickle=True, **kwds):
     """
     Save several arrays into a single file in compressed ``.npz`` format.
 
@@ -708,6 +716,14 @@ def savez_compressed(file, *args, **kwds):
         Arrays to save to the file. Please use keyword arguments (see
         `kwds` below) to assign names to arrays.  Arrays specified as
         args will be named "arr_0", "arr_1", and so on.
+    allow_pickle : bool, optional
+        Allow saving object arrays using Python pickles. Reasons for
+        disallowing pickles include security (loading pickled data can execute
+        arbitrary code) and portability (pickled objects may not be loadable
+        on different Python installations, for example if the stored objects
+        require libraries that are not available, and not all pickled data is
+        compatible between different versions of Python).
+        Default: True
     kwds : Keyword arguments, optional
         Arrays to save to the file. Each array will be saved to the
         output file with its corresponding keyword name.
@@ -750,7 +766,7 @@ def savez_compressed(file, *args, **kwds):
     True
 
     """
-    _savez(file, args, kwds, True)
+    _savez(file, args, kwds, True, allow_pickle=allow_pickle)
 
 
 def _savez(file, args, kwds, compress, allow_pickle=True, pickle_kwargs=None):
@@ -777,17 +793,17 @@ def _savez(file, args, kwds, compress, allow_pickle=True, pickle_kwargs=None):
         compression = zipfile.ZIP_STORED
 
     zipf = zipfile_factory(file, mode="w", compression=compression)
-
-    for key, val in namedict.items():
-        fname = key + '.npy'
-        val = np.asanyarray(val)
-        # always force zip64, gh-10776
-        with zipf.open(fname, 'w', force_zip64=True) as fid:
-            format.write_array(fid, val,
-                               allow_pickle=allow_pickle,
-                               pickle_kwargs=pickle_kwargs)
-
-    zipf.close()
+    try:
+        for key, val in namedict.items():
+            fname = key + '.npy'
+            val = np.asanyarray(val)
+            # always force zip64, gh-10776
+            with zipf.open(fname, 'w', force_zip64=True) as fid:
+                format.write_array(fid, val,
+                                   allow_pickle=allow_pickle,
+                                   pickle_kwargs=pickle_kwargs)
+    finally:
+        zipf.close()
 
 
 def _ensure_ndmin_ndarray_check_param(ndmin):

numpy/lib/_npyio_impl.pyi

@@ -158,12 +158,14 @@ def save(
 def savez(
     file: str | os.PathLike[str] | _SupportsWrite[bytes],
     *args: ArrayLike,
+    allow_pickle: bool = ...,
     **kwds: ArrayLike,
 ) -> None: ...
 
 def savez_compressed(
     file: str | os.PathLike[str] | _SupportsWrite[bytes],
     *args: ArrayLike,
+    allow_pickle: bool = ...,
     **kwds: ArrayLike,
 ) -> None: ...
 

numpy/lib/tests/test_io.py

@@ -2790,6 +2790,7 @@ def test_load_refcount():
         x = np.loadtxt(TextIO("0 1 2 3"), dtype=dt)
         assert_equal(x, np.array([((0, 1), (2, 3))], dtype=dt))
 
+
 def test_load_multiple_arrays_until_eof():
     f = BytesIO()
     np.save(f, 1)
@@ -2799,3 +2800,20 @@ def test_load_multiple_arrays_until_eof():
     assert np.load(f) == 2
     with pytest.raises(EOFError):
         np.load(f)
+
+
+def test_savez_nopickle():
+    obj_array = np.array([1, 'hello'], dtype=object)
+    with temppath(suffix='.npz') as tmp:
+        np.savez(tmp, obj_array)
+
+    with temppath(suffix='.npz') as tmp:
+        with pytest.raises(ValueError, match="Object arrays cannot be saved when.*"):
+            np.savez(tmp, obj_array, allow_pickle=False)
+
+    with temppath(suffix='.npz') as tmp:
+        np.savez_compressed(tmp, obj_array)
+
+    with temppath(suffix='.npz') as tmp:
+        with pytest.raises(ValueError, match="Object arrays cannot be saved when.*"):
+            np.savez_compressed(tmp, obj_array, allow_pickle=False)