make the connection between OWNDATA and having a allocator handle mor…

…e explicit
numpy · seberg · Oct 25, 2021 · Oct 11, 2020 · Oct 16, 2020 · Oct 18, 2020
commit 2945c644d20ee2d5abd0e4062d181ee6bd6058e6
diff --git a/numpy/core/_add_newdocs.py b/numpy/core/_add_newdocs.py
@@ -4690,11 +4690,12 @@
 
 add_newdoc('numpy.core.multiarray', 'get_handler_name',
     """
-    get_handler_name(a: ndarray) -> str
+    get_handler_name(a: ndarray) -> str,None
 
     Return the name of the memory handler used by `a`. If not provided, return
     the name of the current global memory handler that will be used to allocate
-    data for the next `ndarray`.
+    data for the next `ndarray`. May return None if `a` does not own its
+    memory, in which case you can traverse ``a.base`` for a memory handler.
     """)
 
 add_newdoc('numpy.core.multiarray', '_set_madvise_hugepage',

diff --git a/numpy/core/src/multiarray/alloc.c b/numpy/core/src/multiarray/alloc.c
@@ -552,9 +552,10 @@ PyDataMem_SetHandler(PyDataMem_Handler *handler)
 }
 
 /*NUMPY_API
- * Return the PyDataMem_Handler used by the PyArrayObject. If NULL, return
+ * Return the PyDataMem_Handler used by obj. If obj is NULL, return
  * the current global policy that will be used to allocate data
- * for the next PyArrayObject. On failure, return NULL.
+ * for the next PyArrayObject. On failure, return NULL. Can also return NULL
+ * if obj does not own its memory
  */
 NPY_NO_EXPORT const PyDataMem_Handler *
 PyDataMem_GetHandler(PyArrayObject *obj)
@@ -588,20 +589,7 @@ PyDataMem_GetHandler(PyArrayObject *obj)
         return handler;
     }
 #endif
-    /* Try to find a handler */
-    PyArrayObject *base = obj;
-    handler = PyArray_HANDLER(obj);
-    while (handler == NULL) {
-        base = (PyArrayObject*)PyArray_BASE(base);
-        /*
-         * If the base is an array which owns its own data, return its allocator.
-         */
-        if (base == NULL || ! PyArray_Check(base)) {
-            break;
-        }
-        handler = PyArray_HANDLER(base);
-    }
-    return handler;
+    return PyArray_HANDLER(obj);
 }
 
 NPY_NO_EXPORT PyObject *
@@ -617,7 +605,10 @@ get_handler_name(PyObject *NPY_UNUSED(self), PyObject *args)
     }
     const PyDataMem_Handler * mem_handler = PyDataMem_GetHandler((PyArrayObject *)arr);
     if (mem_handler == NULL) {
-        return NULL;
+        if (PyErr_Occurred()) {
+            return NULL;
+        }
+        Py_RETURN_NONE;
     }
     return PyUnicode_FromString(mem_handler->name);
 }
diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
@@ -843,8 +843,7 @@ PyArray_NewFromDescr_int(
         /* The handlers should never be called in this case */
         fa->mem_handler = NULL;
         /*
-         * If data is passed in, this object won't own it by default.
-         * Caller must arrange for this to be reset if truly desired
+         * If data is passed in, this object won't own it.
          */
         fa->flags &= ~NPY_ARRAY_OWNDATA;
     }
@@ -3412,6 +3411,7 @@ array_from_text(PyArray_Descr *dtype, npy_intp num, char const *sep, size_t *nre
         dptr += dtype->elsize;
         if (num < 0 && thisbuf == size) {
             totalbytes += bytes;
+            /* The handler is always valid */
             tmp = PyDataMem_UserRENEW(PyArray_DATA(r), totalbytes,
                                   &PyArray_HANDLER(r)->allocator);
             if (tmp == NULL) {
@@ -3435,6 +3435,7 @@ array_from_text(PyArray_Descr *dtype, npy_intp num, char const *sep, size_t *nre
         const size_t nsize = PyArray_MAX(*nread,1)*dtype->elsize;
 
         if (nsize != 0) {
+            /* The handler is always valid */
             tmp = PyDataMem_UserRENEW(PyArray_DATA(r), nsize,
                                   &PyArray_HANDLER(r)->allocator);
             if (tmp == NULL) {
@@ -3541,6 +3542,7 @@ PyArray_FromFile(FILE *fp, PyArray_Descr *dtype, npy_intp num, char *sep)
         const size_t nsize = PyArray_MAX(nread,1) * dtype->elsize;
         char *tmp;
 
+        /* The handler is always valid */
         if((tmp = PyDataMem_UserRENEW(PyArray_DATA(ret), nsize,
                                      &PyArray_HANDLER(ret)->allocator)) == NULL) {
             Py_DECREF(dtype);
@@ -3826,6 +3828,7 @@ PyArray_FromIter(PyObject *obj, PyArray_Descr *dtype, npy_intp count)
             */
             elcount = (i >> 1) + (i < 4 ? 4 : 2) + i;
             if (!npy_mul_with_overflow_intp(&nbytes, elcount, elsize)) {
+                /* The handler is always valid */
                 new_data = PyDataMem_UserRENEW(PyArray_DATA(ret), nbytes,
                                   &PyArray_HANDLER(ret)->allocator);
             }
@@ -3868,6 +3871,7 @@ PyArray_FromIter(PyObject *obj, PyArray_Descr *dtype, npy_intp count)
         /* The size cannot be zero for realloc. */
         goto done;
     }
+    /* The handler is always valid */
     new_data = PyDataMem_UserRENEW(PyArray_DATA(ret), i * elsize,
                                    &PyArray_HANDLER(ret)->allocator);
     if (new_data == NULL) {

diff --git a/numpy/core/src/multiarray/getset.c b/numpy/core/src/multiarray/getset.c
@@ -393,8 +393,14 @@ array_data_set(PyArrayObject *self, PyObject *op, void *NPY_UNUSED(ignored))
             PyArray_Descr *dtype = PyArray_DESCR(self);
             nbytes = dtype->elsize ? dtype->elsize : 1;
         }
-        PyDataMem_UserFREE(PyArray_DATA(self), nbytes,
-                           &PyArray_HANDLER(self)->allocator);
+        PyDataMem_Handler *handler = PyArray_HANDLER(self);
+        if (handler == NULL) {
+            /* This can happen if someone arbitrarily sets NPY_ARRAY_OWNDATA */
+            PyErr_SetString(PyExc_RuntimeError,
+                            "no memory handler found but OWNDATA flag set");
+            return -1;
+        }
+        PyDataMem_UserFREE(PyArray_DATA(self), nbytes, &handler->allocator);
     }
     if (PyArray_BASE(self)) {
         if ((PyArray_FLAGS(self) & NPY_ARRAY_WRITEBACKIFCOPY) ||

diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c
@@ -2053,8 +2
10000
053,14 @@ array_setstate(PyArrayObject *self, PyObject *args)
          * Allocation will never be 0, see comment in ctors.c
          * line 820
          */
-        PyDataMem_UserFREE(PyArray_DATA(self), n_tofree,
-                           &PyArray_HANDLER(self)->allocator);
+        PyDataMem_Handler *handler = PyArray_HANDLER(self);
+        if (handler == NULL) {
+            /* This can happen if someone arbitrarily sets NPY_ARRAY_OWNDATA */
+            PyErr_SetString(PyExc_RuntimeError,
+                            "no memory handler found but OWNDATA flag set");
+            return NULL;
+        }
+        PyDataMem_UserFREE(PyArray_DATA(self), n_tofree, &handler->allocator);
         PyArray_CLEARFLAGS(self, NPY_ARRAY_OWNDATA);
     }
     Py_XDECREF(PyArray_BASE(self));

diff --git a/numpy/core/src/multiarray/shape.c b/numpy/core/src/multiarray/shape.c
@@ -120,9 +120,16 @@ PyArray_Resize(PyArrayObject *self, PyArray_Dims *newshape, int refcheck,
         }
 
         /* Reallocate space if needed - allocating 0 is forbidden */
+        PyDataMem_Handler *handler = PyArray_HANDLER(self);
+        if (handler == NULL) {
+            /* This can happen if someone arbitrarily sets NPY_ARRAY_OWNDATA */
+            PyErr_SetString(PyExc_RuntimeError,
+                            "no memory handler found but OWNDATA flag set");
+            return NULL;
+        }
         new_data = PyDataMem_UserRENEW(PyArray_DATA(self),
                                        newnbytes == 0 ? elsize : newnbytes,
-                                       &PyArray_HANDLER(self)->allocator);
+                                       &handler->allocator);
         if (new_data == NULL) {
             PyErr_SetString(PyExc_MemoryError,
                     "cannot allocate memory for array");

diff --git a/numpy/core/tests/test_mem_policy.py b/numpy/core/tests/test_mem_policy.py
@@ -151,12 +151,14 @@ def test_set_policy(get_module):
     orig_policy_name = np.core.multiarray.get_handler_name()
 
     a = np.arange(10).reshape((2, 5))  # a doesn't own its own data
-    assert np.core.multiarray.get_handler_name(a) == orig_policy_name
+    assert np.core.multiarray.get_handler_name(a) == None
+    assert np.core.multiarray.get_handler_name(a.base) == orig_policy_name
 
     orig_policy = get_module.set_secret_data_policy()
 
     b = np.arange(10).reshape((2, 5))  # b doesn't own its own data
-    assert np.core.multiarray.get_handler_name(b) == 'secret_data_allocator'
+    assert np.core.multiarray.get_handler_name(b) == None
+    assert np.core.multiarray.get_handler_name(b.base) == 'secret_data_allocator'
 
     if orig_policy_name == 'default_allocator':
         get_module.set_old_policy(None)  # tests PyDataMem_SetHandler(NULL)
@@ -171,16 +173,17 @@ def test_policy_propagation(get_module):
     class MyArr(np.ndarray):
         pass
 
+    # The memory policy goes hand-in-hand with flags.owndata
     orig_policy_name = np.core.multiarray.get_handler_name()
-    a = np.arange(10).view(MyArr).reshape((2, 5))  # a doesn't own its own data
-    assert np.core.multiarray.get_handler_name(a) == orig_policy_name
+    a = np.arange(10).view(MyArr).reshape((2, 5))
+    assert np.core.multiarray.get_handler_name(a) == None
+    assert a.flags.owndata == False
 
-    orig_policy = get_module.set_secret_data_policy()
-    secret_policy_name = np.core.multiarray.get_handler_name()
-    b = np.arange(10).view(MyArr).reshape((2, 5))  # b doesn't own its own data
-    assert np.core.multiarray.get_handler_name(b) == secret_policy_name
-    get_module.set_old_policy(orig_policy)
+    assert np.core.multiarray.get_handler_name(a.base) == None
+    assert a.base.flags.owndata == False
 
+    assert np.core.multiarray.get_handler_name(a.base.base) == orig_policy_name
+    assert a.base.base.flags.owndata == True
 
 async def concurrent_context1(get_module, orig_policy_name, event):
     if orig_policy_name == 'default_allocator':

diff --git a/numpy/core/tests/test_nditer.py b/numpy/core/tests/test_nditer.py
@@ -3149,6 +3149,7 @@ def test_partial_iteration_cleanup(in_dtype, buf_dtype, steps):
     # Note that resetting does not free references
     del it
     break_cycles()
+    break_cycles()
     assert count == sys.getrefcount(value)
 
     # Repeat the test with `iternext`
@@ -3159,6 +3160,7 @@ def test_partial_iteration_cleanup(in_dtype, buf_dtype, steps):
 
     del it  # should ensure cleanup
     break_cycles()
+    break_cycles()
     assert count == sys.getrefcount(value)