ENH: fix and test for blindly taking ownership of data

mattip · mattip · commit e81eff4a2844 · 2021-10-25T18:53:45.000+03:00
diff --git a/numpy/core/src/multiarray/arrayobject.c b/numpy/core/src/multiarray/arrayobject.c
@@ -501,8 +501,17 @@ array_dealloc(PyArrayObject *self)
         if (nbytes == 0) {
             nbytes = fa->descr->elsize ? fa->descr->elsize : 1;
         }
-        PyDataMem_UserFREE(fa->data, nbytes, fa->mem_handler);
-        Py_DECREF(fa->mem_handler);
+        if (fa->mem_handler == NULL) {
+            char const * msg = "Trying to dealloc data, but a memory policy "
+                "is not set. If you take ownership of the data, you must "
+                "also set a memory policy.";
+            WARN_IN_DEALLOC(PyExc_RuntimeWarning, msg);
+            // Guess at malloc/free ???
+            free(fa->data);
+        } else {
+            PyDataMem_UserFREE(fa->data, nbytes, fa->mem_handler);
+            Py_DECREF(fa->mem_handler);
+        }
     }
 
     /* must match allocation in PyArray_NewFromDescr */
diff --git a/numpy/core/tests/test_mem_policy.py b/numpy/core/tests/test_mem_policy.py
@@ -2,7 +2,8 @@
 import pytest
 import numpy as np
 import threading
-from numpy.testing import extbuild
+import warnings
+from numpy.testing import extbuild, assert_warns
 import sys
 
 
@@ -40,6 +41,25 @@ def get_module(tmp_path):
              Py_DECREF(old);
              Py_RETURN_NONE;
          """),
+        ("get_array", "METH_NOARGS", """
+            char *buf = (char *)malloc(20);
+            npy_intp dims[1];
+            dims[0] = 20;
+            PyArray_Descr *descr =  PyArray_DescrNewFromType(NPY_UINT8);
+            return PyArray_NewFromDescr(&PyArray_Type, descr, 1, dims, NULL,
+                                        buf, NPY_ARRAY_WRITEABLE, NULL);
+         """),
+        ("set_own", "METH_O", """
+            if (!PyArray_Check(args)) {
+                PyErr_SetString(PyExc_ValueError,
+                             "need an ndarray");
+                return NULL;
+            }
+            PyArray_ENABLEFLAGS((PyArrayObject*)args, NPY_ARRAY_OWNDATA);
+            // Maybe try this too?
+            // PyArray_BASE(PyArrayObject *)args) = NULL;
+            Py_RETURN_NONE;
+         """),
     ]
     prologue = '''
         #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
@@ -305,3 +325,14 @@ def test_new_policy(get_module):
 
     c = np.arange(10)
     assert np.core.multiarray.get_handler_name(c) == orig_policy_name
+
+def test_switch_owner(get_module):
+    a = get_module.get_array()
+    assert np.core.multiarray.get_handler_name(a) is None
+    get_module.set_own(a)
+    with warnings.catch_warnings():
+        warnings.filterwarnings('always')
+        # The policy should be NULL, so we have to assume we can call "free"
+        with assert_warns(RuntimeWarning) as w:
+            del a
+        print(w)
