Merge pull request #12011 from pierreglaser/implement-reduce-ex

ENH: implementation of array_reduce_ex

1 files changed, 131 insertions, 0 deletions

diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c
index cdbd0d6ae..23b0bfd24 100644
--- a/numpy/core/src/multiarray/methods.c
+++ b/numpy/core/src/multiarray/methods.c
@@ -1619,6 +1619,8 @@ array_reduce(PyArrayObject *self, PyObject *NPY_UNUSED(args))
 
        Notice because Python does not describe a mechanism to write
        raw data to the pickle, this performs a copy to a string first
+       This issue is now adressed in protocol 5, where a buffer is serialized
+       instead of a string,
     */
 
     state = PyTuple_New(5);
@@ -1652,6 +1654,132 @@ array_reduce(PyArrayObject *self, PyObject *NPY_UNUSED(args))
 }
 
 static PyObject *
+array_reduce_ex(PyArrayObject *self, PyObject *args)
+{
+    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 *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;
+            }
+
+            /* 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;
+            }
+#else
+            PyErr_SetString(PyExc_ValueError, "pickle protocol 5 is not available "
+                                               "for python versions < 3.6");
+            return NULL;
+#endif
+            if (pickle_module == NULL){
+                return NULL;
+            }
+
+            pickle_class = PyObject_GetAttrString(pickle_module,
+                                                  "PickleBuffer");
+
+            class_args_tuple = PyTuple_New(1);
+            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);
+                PyTuple_SET_ITEM(class_args_tuple, 0, transposed_array);
+            }
+            else {
+                order = 'C';
+                PyTuple_SET_ITEM(class_args_tuple, 0, (PyObject *)self);
+                Py_INCREF(self);
+            }
+
+            class_args = Py_BuildValue("O", class_args_tuple);
+
+            buffer = PyObject_CallObject(pickle_class, class_args);
+
+            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);
+
+            Py_INCREF(descr);
+
+            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));
+
+            PyTuple_SET_ITEM(ret, 0, from_buffer_func);
+            PyTuple_SET_ITEM(ret, 1, buffer_tuple);
+
+            return ret;
+        }
+        else {
+            PyErr_Format(PyExc_ValueError,
+                         "cannot call __reduce_ex__ with protocol >= %d",
+                         5);
+            return NULL;
+        }
+    }
+    else {
+        return NULL;
+    }
+
+}
+
+static PyObject *
 array_setstate(PyArrayObject *self, PyObject *args)
 {
     PyObject *shape;
@@ -2524,6 +2652,9 @@ NPY_NO_EXPORT PyMethodDef array_methods[] = {
     {"__reduce__",
         (PyCFunction) array_reduce,
         METH_VARARGS, NULL},
+    {"__reduce_ex__",
+        (PyCFunction) array_reduce_ex,
+        METH_VARARGS, NULL},
     {"__setstate__",
         (PyCFunction) array_setstate,
         METH_VARARGS, NULL},