Merge pull request #24906 from mtsokol/backport-_core-stubs

MAINT: Backport `numpy._core` stubs. Remove `NumpyUnpickler`

doc/release/upcoming_changes/24906.new_feature.rst

@@ -0,0 +1,6 @@
+`numpy._core` submodules' stubs
+-------------------------------
+
+`numpy._core` submodules' stubs were added
+to provide a stable way for loading pickled arrays,
+created with NumPy 2.0, with Numpy 1.26.

doc/source/reference/routines.io.rst

@@ -14,7 +14,6 @@ NumPy binary files (NPY, NPZ)
    save
    savez
    savez_compressed
-   lib.format.NumpyUnpickler
 
 The format of these binary file types is documented in
 :py:mod:`numpy.lib.format`

doc/source/user/how-to-io.rst

@@ -319,10 +319,8 @@ Use :func:`numpy.save` and :func:`numpy.load`.  Set ``allow_pickle=False``,
 unless the array dtype includes Python objects, in which case pickling is
 required.
 
-:func:`numpy.load` also supports unpickling files created with NumPy 2.0.
-If you try to unpickle a 2.0 pickled array directly, you will get
-an exception. Use :class:`numpy.lib.format.NumpyUnpickler` for
-unpickling these files.
+NumPy 1.26 also supports unpickling files created with NumPy 2.0, either
+via :func:`numpy.load` or the pickle module directly.
 
 Convert from a pandas DataFrame to a NumPy array
 ================================================

numpy/_core/__init__.py

@@ -0,0 +1,4 @@
+"""
+This private module only contains stubs for interoperability with
+NumPy 2.0 pickled arrays. It may not be used by the end user.
+"""

numpy/_core/_dtype.py

@@ -0,0 +1,6 @@
+from numpy.core import _dtype
+
+_globals = globals()
+
+for item in _dtype.__dir__():
+    _globals[item] = getattr(_dtype, item)

numpy/_core/_dtype_ctypes.py

@@ -0,0 +1,6 @@
+from numpy.core import _dtype_ctypes
+
+_globals = globals()
+
+for item in _dtype_ctypes.__dir__():
+    _globals[item] = getattr(_dtype_ctypes, item)

numpy/_core/_internal.py

@@ -0,0 +1,6 @@
+from numpy.core import _internal
+
+_globals = globals()
+
+for item in _internal.__dir__():
+    _globals[item] = getattr(_internal, item)

numpy/_core/_multiarray_umath.py

@@ -0,0 +1,6 @@
+from numpy.core import _multiarray_umath
+
+_globals = globals()
+
+for item in _multiarray_umath.__dir__():
+    _globals[item] = getattr(_multiarray_umath, item)

numpy/_core/multiarray.py

@@ -0,0 +1,6 @@
+from numpy.core import multiarray
+
+_globals = globals()
+
+for item in multiarray.__dir__():
+    _globals[item] = getattr(multiarray, item)

numpy/_core/umath.py

@@ -0,0 +1,6 @@
+from numpy.core import umath
+
+_globals = globals()
+
+for item in umath.__dir__():
+    _globals[item] = getattr(umath, item)

numpy/core/tests/test_numpy_2_0_compat.py

@@ -0,0 +1,48 @@
+from os import path
+import pickle
+
+import numpy as np
+
+
+class TestNumPy2Compatibility:
+
+    data_dir = path.join(path.dirname(__file__), "data")
+    filename = path.join(data_dir, "numpy_2_0_array.pkl")
+
+    def test_importable__core_stubs(self):
+        """
+        Checks if stubs for `numpy._core` are importable.
+        """
+        from numpy._core.multiarray import _reconstruct
+        from numpy._core.umath import cos
+        from numpy._core._multiarray_umath import exp
+        from numpy._core._internal import ndarray
+        from numpy._core._dtype import _construction_repr
+        from numpy._core._dtype_ctypes import dtype_from_ctypes_type
+
+    def test_unpickle_numpy_2_0_file(self):
+        """
+        Checks that NumPy 1.26 and pickle is able to load pickles
+        created with NumPy 2.0 without errors/warnings.
+        """
+        with open(self.filename, mode="rb") as file:
+            content = file.read()
+
+        # Let's make sure that the pickle object we're loading
+        # was built with NumPy 2.0.
+        assert b"numpy._core.multiarray" in content
+
+        arr = pickle.loads(content, encoding="latin1")
+
+        assert isinstance(arr, np.ndarray)
+        assert arr.shape == (73,) and arr.dtype == np.float64
+
+    def test_numpy_load_numpy_2_0_file(self):
+        """
+        Checks that `numpy.load` for NumPy 1.26 is able to load pickles
+        created with NumPy 2.0 without errors/warnings.
+        """
+        arr = np.load(self.filename, encoding="latin1", allow_pickle=True)
+
+        assert isinstance(arr, np.ndarray)
+        assert arr.shape == (73,) and arr.dtype == np.float64

numpy/lib/format.py

@@ -169,7 +169,7 @@
     )
 
 
-__all__ = ["NumpyUnpickler"]
+__all__ = []
 
 
 EXPECTED_KEYS = {'descr', 'fortran_order', 'shape'}
@@ -797,7 +797,7 @@ def read_array(fp, allow_pickle=False, pickle_kwargs=None, *,
         if pickle_kwargs is None:
             pickle_kwargs = {}
         try:
-            array = NumpyUnpickler(fp, **pickle_kwargs).load()
+            array = pickle.load(fp, **pickle_kwargs)
         except UnicodeError as err:
             # Friendlier error message
             raise UnicodeError("Unpickling a python object failed: %r\n"
@@ -974,15 +974,3 @@ def _read_bytes(fp, size, error_template="ran out of data"):
         raise ValueError(msg % (error_template, size, len(data)))
     else:
         return data
-
-
-class NumpyUnpickler(pickle.Unpickler):
-    """
-    A thin wrapper for :py:class:`pickle.Unpickler` that
-    allows to load 2.0 array pickles with numpy 1.26.
-    """
-
-    def find_class(self, module: str, name: str) -> object:
-        if module.startswith("numpy._core"):
-            module = module.replace("_core", "core", 1)
-        return pickle.Unpickler.find_class(self, module, name)

numpy/lib/format.pyi

@@ -1,4 +1,3 @@
-import pickle
 from typing import Any, Literal, Final
 
 __all__: list[str]
@@ -21,5 +20,3 @@ def read_array_header_2_0(fp): ...
 def write_array(fp, array, version=..., allow_pickle=..., pickle_kwargs=...): ...
 def read_array(fp, allow_pickle=..., pickle_kwargs=...): ...
 def open_memmap(filename, mode=..., dtype=..., shape=..., fortran_order=..., version=...): ...
-
-class NumpyUnpickler(pickle.Unpickler): ...

numpy/lib/npyio.py

@@ -462,7 +462,7 @@ def load(file, mmap_mode=None, allow_pickle=False, fix_imports=True,
                 raise ValueError("Cannot load file containing pickled data "
                                  "when allow_pickle=False")
             try:
-                return format.NumpyUnpickler(fid, **pickle_kwargs).load()
+                return pickle.load(fid, **pickle_kwargs)
             except Exception as e:
                 raise pickle.UnpicklingError(
                     f"Failed to interpret file {file!r} as a pickle") from e

numpy/meson.build

@@ -246,6 +246,7 @@ pure_subdirs = [
   '_pyinstaller',
   '_typing',
   '_utils',
+  '_core',
   'array_api',
   'compat',
   'doc',

numpy/setup.py

@@ -7,6 +7,7 @@ def configuration(parent_package='',top_path=None):
     config.add_subpackage('array_api')
     config.add_subpackage('compat')
     config.add_subpackage('core')
+    config.add_subpackage('_core')
     config.add_subpackage('distutils')
     config.add_subpackage('doc')
     config.add_subpackage('f2py')