diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c index 6005b9751bf4..b519f75fc851 100644 --- a/numpy/core/src/multiarray/methods.c +++ b/numpy/core/src/multiarray/methods.c @@ -1723,129 +1723,150 @@ array_reduce(PyArrayObject *self, PyObject *NPY_UNUSED(args)) } static PyObject * -array_reduce_ex(PyArrayObject *self, PyObject *args) +array_reduce_ex_regular(PyArrayObject *self, int protocol) { - int protocol; - PyObject *ret = NULL, *numeric_mod = NULL, *from_buffer_func = NULL; - PyObject *buffer_tuple = NULL, *pickle_module = NULL, *pickle_class = NULL; - PyObject *class_args = NULL, *class_args_tuple = NULL, *unused = NULL; PyObject *subclass_array_reduce = NULL; + PyObject *ret; + + /* We do not call array_reduce directly but instead lookup and call + * the __reduce__ method to make sure that it's possible to customize + * pickling in sub-classes. */ + subclass_array_reduce = PyObject_GetAttrString((PyObject *)self, + "__reduce__"); + if (subclass_array_reduce == NULL) { + return NULL; + } + ret = PyObject_CallObject(subclass_array_reduce, NULL); + Py_DECREF(subclass_array_reduce); + return ret; +} + +static PyObject * +array_reduce_ex_picklebuffer(PyArrayObject *self, int protocol) +{ + PyObject *numeric_mod = NULL, *from_buffer_func = NULL; + PyObject *pickle_module = NULL, *picklebuf_class = NULL; + PyObject *picklebuf_args = NULL; PyObject *buffer = NULL, *transposed_array = NULL; PyArray_Descr *descr = NULL; char order; - if (PyArg_ParseTuple(args, "i", &protocol)){ - descr = PyArray_DESCR(self); - if ((protocol < 5) || - (!PyArray_IS_C_CONTIGUOUS((PyArrayObject*)self) && - !PyArray_IS_F_CONTIGUOUS((PyArrayObject*)self)) || - PyDataType_FLAGCHK(descr, NPY_ITEM_HASOBJECT) || - (PyType_IsSubtype(((PyObject*)self)->ob_type, &PyArray_Type) && - ((PyObject*)self)->ob_type != &PyArray_Type) || - PyDataType_ISUNSIZED(descr)) { - /* The PickleBuffer class from version 5 of the pickle protocol - * can only be used for arrays backed by a contiguous data buffer. - * For all other cases we fallback to the generic array_reduce - * method that involves using a temporary bytes allocation. However - * we do not call array_reduce directly but instead lookup and call - * the __reduce__ method to make sure that it's possible customize - * pickling in sub-classes. */ - subclass_array_reduce = PyObject_GetAttrString((PyObject *)self, - "__reduce__"); - return PyObject_CallObject(subclass_array_reduce, unused); - } - else if (protocol == 5){ - ret = PyTuple_New(2); - - if (ret == NULL) { - return NULL; - } + descr = PyArray_DESCR(self); - /* if the python version is below 3.8, the pickle module does not provide - * built-in support for protocol 5. We try importing the pickle5 - * backport instead */ + /* if the python version is below 3.8, the pickle module does not provide + * built-in support for protocol 5. We try importing the pickle5 + * backport instead */ #if PY_VERSION_HEX >= 0x03080000 - pickle_module = PyImport_ImportModule("pickle"); -#elif PY_VERSION_HEX < 0x03080000 && PY_VERSION_HEX >= 0x03060000 - pickle_module = PyImport_ImportModule("pickle5"); - if (pickle_module == NULL){ - /* for protocol 5, raise a clear ImportError if pickle5 is not found - */ - PyErr_SetString(PyExc_ImportError, "Using pickle protocol 5 " - "requires the pickle5 module for python versions >=3.6 " - "and <3.8"); - return NULL; - } + /* we expect protocol 5 to be available in Python 3.8 */ + pickle_module = PyImport_ImportModule("pickle"); +#elif PY_VERSION_HEX >= 0x03060000 + pickle_module = PyImport_ImportModule("pickle5"); + if (pickle_module == NULL) { + /* for protocol 5, raise a clear ImportError if pickle5 is not found + */ + PyErr_SetString(PyExc_ImportError, "Using pickle protocol 5 " + "requires the pickle5 module for Python >=3.6 and <3.8"); + return NULL; + } #else - PyErr_SetString(PyExc_ValueError, "pickle protocol 5 is not available " - "for python versions < 3.6"); - return NULL; + PyErr_SetString(PyExc_ValueError, "pickle protocol 5 is not available " + "for Python < 3.6"); + return NULL; #endif - if (pickle_module == NULL){ - return NULL; - } - - pickle_class = PyObject_GetAttrString(pickle_module, - "PickleBuffer"); + if (pickle_module == NULL){ + return NULL; + } + picklebuf_class = PyObject_GetAttrString(pickle_module, "PickleBuffer"); + Py_DECREF(pickle_module); + if (picklebuf_class == NULL) { + return NULL; + } - class_args_tuple = PyTuple_New(1); - if (!PyArray_IS_C_CONTIGUOUS((PyArrayObject*)self) && - PyArray_IS_F_CONTIGUOUS((PyArrayObject*)self)){ + /* Construct a PickleBuffer of the array */ - /* if the array if Fortran-contiguous and not C-contiguous, - * the PickleBuffer instance will hold a view on the transpose - * of the initial array, that is C-contiguous. */ - order = 'F'; - transposed_array = PyArray_Transpose((PyArrayObject*)self, NULL); - PyTuple_SET_ITEM(class_args_tuple, 0, transposed_array); - } - else { - order = 'C'; - PyTuple_SET_ITEM(class_args_tuple, 0, (PyObject *)self); - Py_INCREF(self); - } + if (!PyArray_IS_C_CONTIGUOUS((PyArrayObject*) self) && + PyArray_IS_F_CONTIGUOUS((PyArrayObject*) self)) { + /* if the array if Fortran-contiguous and not C-contiguous, + * the PickleBuffer instance will hold a view on the transpose + * of the initial array, that is C-contiguous. */ + order = 'F'; + transposed_array = PyArray_Transpose((PyArrayObject*)self, NULL); + picklebuf_args = Py_BuildValue("(N)", transposed_array); + } + else { + order = 'C'; + picklebuf_args = Py_BuildValue("(O)", self); + } + if (picklebuf_args == NULL) { + Py_DECREF(picklebuf_class); + return NULL; + } - class_args = Py_BuildValue("O", class_args_tuple); + buffer = PyObject_CallObject(picklebuf_class, picklebuf_args); + Py_DECREF(picklebuf_class); + Py_DECREF(picklebuf_args); + if (buffer == NULL) { + /* Some arrays may refuse to export a buffer, in which case + * just fall back on regular __reduce_ex__ implementation + * (gh-12745). + */ + PyErr_Clear(); + return array_reduce_ex_regular(self, protocol); + } - buffer = PyObject_CallObject(pickle_class, class_args); + /* Get the _frombuffer() function for reconstruction */ - numeric_mod = PyImport_ImportModule("numpy.core.numeric"); - if (numeric_mod == NULL) { - Py_DECREF(ret); - return NULL; - } - from_buffer_func = PyObject_GetAttrString(numeric_mod, - "_frombuffer"); - Py_DECREF(numeric_mod); + numeric_mod = PyImport_ImportModule("numpy.core.numeric"); + if (numeric_mod == NULL) { + Py_DECREF(buffer); + return NULL; + } + from_buffer_func = PyObject_GetAttrString(numeric_mod, + "_frombuffer"); + Py_DECREF(numeric_mod); + if (from_buffer_func == NULL) { + Py_DECREF(buffer); + return NULL; + } - Py_INCREF(descr); + return Py_BuildValue("N(NONN)", + from_buffer_func, buffer, (PyObject *)descr, + PyObject_GetAttrString((PyObject *)self, "shape"), + PyUnicode_FromStringAndSize(&order, 1)); +} - buffer_tuple = PyTuple_New(4); - PyTuple_SET_ITEM(buffer_tuple, 0, buffer); - PyTuple_SET_ITEM(buffer_tuple, 1, (PyObject *)descr); - PyTuple_SET_ITEM(buffer_tuple, 2, - PyObject_GetAttrString((PyObject *)self, - "shape")); - PyTuple_SET_ITEM(buffer_tuple, 3, - PyUnicode_FromStringAndSize(&order, - (Py_ssize_t)1)); +static PyObject * +array_reduce_ex(PyArrayObject *self, PyObject *args) +{ + int protocol; + PyArray_Descr *descr = NULL; - PyTuple_SET_ITEM(ret, 0, from_buffer_func); - PyTuple_SET_ITEM(ret, 1, buffer_tuple); + if (!PyArg_ParseTuple(args, "i", &protocol)) { + return NULL; + } - return ret; - } - else { - PyErr_Format(PyExc_ValueError, - "cannot call __reduce_ex__ with protocol >= %d", - 5); - return NULL; - } + descr = PyArray_DESCR(self); + if ((protocol < 5) || + (!PyArray_IS_C_CONTIGUOUS((PyArrayObject*)self) && + !PyArray_IS_F_CONTIGUOUS((PyArrayObject*)self)) || + PyDataType_FLAGCHK(descr, NPY_ITEM_HASOBJECT) || + (PyType_IsSubtype(((PyObject*)self)->ob_type, &PyArray_Type) && + ((PyObject*)self)->ob_type != &PyArray_Type) || + PyDataType_ISUNSIZED(descr)) { + /* The PickleBuffer class from version 5 of the pickle protocol + * can only be used for arrays backed by a contiguous data buffer. + * For all other cases we fallback to the generic array_reduce + * method that involves using a temporary bytes allocation. */ + return array_reduce_ex_regular(self, protocol); + } + else if (protocol == 5) { + return array_reduce_ex_picklebuffer(self, protocol); } else { + PyErr_Format(PyExc_ValueError, + "__reduce_ex__ called with protocol > 5"); return NULL; } - } static PyObject * diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py index 241f8e48aa26..dff0facc6ad5 100644 --- a/numpy/core/tests/test_multiarray.py +++ b/numpy/core/tests/test_multiarray.py @@ -3769,10 +3769,16 @@ def test_roundtrip(self): ('c', float)]) ] + refs = [weakref.ref(a) for a in DATA] for a in DATA: assert_equal( a, pickle.loads(pickle.dumps(a, protocol=proto)), err_msg="%r" % a) + del a, DATA, carray + gc.collect() + # check for reference leaks (gh-12793) + for ref in refs: + assert ref() is None def _loads(self, obj): if sys.version_info[0] >= 3: diff --git a/numpy/core/tests/test_regression.py b/numpy/core/tests/test_regression.py index 17c48987f992..d809c285e24e 100644 --- a/numpy/core/tests/test_regression.py +++ b/numpy/core/tests/test_regression.py @@ -2423,3 +2423,11 @@ def test_eff1d_casting(self): assert_equal(res, [-99, 1, 2, 3, -7, 88, 99]) assert_raises(ValueError, np.ediff1d, x, to_begin=(1<<20)) assert_raises(ValueError, np.ediff1d, x, to_end=(1<<20)) + + def test_pickle_datetime64_array(self): + # gh-12745 (would fail with pickle5 installed) + d = np.datetime64('2015-07-04 12:59:59.50', 'ns') + arr = np.array([d]) + for proto in range(2, pickle.HIGHEST_PROTOCOL + 1): + dumped = pickle.dumps(arr, protocol=proto) + assert_equal(pickle.loads(dumped), arr)