diff options
| -rw-r--r-- | numpy/core/__init__.py | 18 | ||||
| -rw-r--r-- | numpy/core/src/umath/_umath_tests.c.src | 9 | ||||
| -rw-r--r-- | numpy/core/tests/test_ufunc.py | 13 |
3 files changed, 32 insertions, 8 deletions
diff --git a/numpy/core/__init__.py b/numpy/core/__init__.py index c77885954..a0769cc89 100644 --- a/numpy/core/__init__.py +++ b/numpy/core/__init__.py @@ -113,10 +113,9 @@ __all__ += getlimits.__all__ __all__ += shape_base.__all__ __all__ += einsumfunc.__all__ -# Make it possible so that ufuncs can be pickled -# Here are the loading and unloading functions -# The name numpy.core._ufunc_reconstruct must be -# available for unpickling to work. +# We used to use `np.core._ufunc_reconstruct` to unpickle. This is unnecessary, +# but old pickles saved before 1.20 will be using it, and there is no reason +# to break loading them. def _ufunc_reconstruct(module, name): # The `fromlist` kwarg is required to ensure that `mod` points to the # inner-most module rather than the parent package when module name is @@ -126,14 +125,17 @@ def _ufunc_reconstruct(module, name): return getattr(mod, name) def _ufunc_reduce(func): - from pickle import whichmodule - name = func.__name__ - return _ufunc_reconstruct, (whichmodule(func, name), name) + # Report the `__name__`. pickle will try to find the module. Note that + # pickle supports for this `__name__` to be a `__qualname__`. It may + # make sense to add a `__qualname__` to ufuncs, to allow this more + # explicitly (Numba has ufuncs as attributes). + # See also: https://github.com/dask/distributed/issues/3450 + return func.__name__ import copyreg -copyreg.pickle(ufunc, _ufunc_reduce, _ufunc_reconstruct) +copyreg.pickle(ufunc, _ufunc_reduce) # Unclutter namespace (must keep _ufunc_reconstruct for unpickling) del copyreg del _ufunc_reduce diff --git a/numpy/core/src/umath/_umath_tests.c.src b/numpy/core/src/umath/_umath_tests.c.src index 660c296d6..750fbeb92 100644 --- a/numpy/core/src/umath/_umath_tests.c.src +++ b/numpy/core/src/umath/_umath_tests.c.src @@ -461,6 +461,15 @@ addUfuncs(PyObject *dictionary) { PyDict_SetItemString(dictionary, "cross1d", f); Py_DECREF(f); + f = PyUFunc_FromFuncAndDataAndSignature(NULL, NULL, + NULL, 0, 0, 0, PyUFunc_None, "_pickleable_module_global.ufunc", + "A dotted name for pickle testing, does nothing.", 0, NULL); + if (f == NULL) { + return -1; + } + PyDict_SetItemString(dictionary, "_pickleable_module_global_ufunc", f); + Py_DECREF(f); + return 0; } diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py index 9eaa1a977..0e9760853 100644 --- a/numpy/core/tests/test_ufunc.py +++ b/numpy/core/tests/test_ufunc.py @@ -178,6 +178,10 @@ class TestUfuncGenericLoops: assert_array_equal(res_num.astype("O"), res_obj) +def _pickleable_module_global(): + pass + + class TestUfunc: def test_pickle(self): for proto in range(2, pickle.HIGHEST_PROTOCOL + 1): @@ -195,6 +199,15 @@ class TestUfunc: b"(S'numpy.core.umath'\np1\nS'cos'\np2\ntp3\nRp4\n.") assert_(pickle.loads(astring) is np.cos) + def test_pickle_name_is_qualname(self): + # This tests that a simplification of our ufunc pickle code will + # lead to allowing qualnames as names. Future ufuncs should + # possible add a specific qualname, or a hook into pickling instead + # (dask+numba may benefit). + _pickleable_module_global.ufunc = umt._pickleable_module_global_ufunc + obj = pickle.loads(pickle.dumps(_pickleable_module_global.ufunc)) + assert obj is umt._pickleable_module_global_ufunc + def test_reduceat_shifting_sum(self): L = 6 x = np.arange(L) |