ENH: expose PyDataMem_NEW/FREE/RENEW as numpy API functions with an event hook.

Moves PyDataMem_NEW/FREE/RENEW to the external API.
Fixes PyDataMem_NEW/RENEW to return void* instead of char*.
Replaces PyDataMem_NEW/FREE with NpySortArray_malloc/free in sort.c.src
  (should be reverted if npysort is moved to be part of multiarraymodule).
Adds PyDataMem_SetEventHook which takes a (PyDataMem_EventHookFunc *) as an argument,
  with signature:
       void hook(void *old, void *new, size_t size).
  When not NULL, hook will be called at the end of each PyDataMem_NEW/FREE/RENEW:
   result = PyDataMem_NEW(size)        -> (*hook(NULL, result, size)
   PyDataMem_FREE(ptr)                 -> (*hook(ptr, NULL, 0)
   result = PyDataMem_RENEW(ptr, size) -> (*hook)(ptr, result, size)
Adds tests in multiarray_tests.c.src, driven by tests/test_multiarray.py.

7 files changed, 161 insertions, 13 deletions

diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
index 15b868e23..b4a5ff90c 100644
--- a/numpy/core/code_generators/numpy_api.py
+++ b/numpy/core/code_generators/numpy_api.py
@@ -346,6 +346,10 @@ multiarray_funcs_api = {
     'PyArray_OutputAllowNAConverter':       306,
     'PyArray_FailUnlessWriteable':          307,
     'PyArray_SetUpdateIfCopyBase':          308,
+    'PyDataMem_NEW':                        309,
+    'PyDataMem_FREE':                       310,
+    'PyDataMem_RENEW':                      311,
+    'PyDataMem_SetEventHook':               312,
 }
 
 ufunc_types_api = {
diff --git a/numpy/core/include/numpy/ndarraytypes.h b/numpy/core/include/numpy/ndarraytypes.h
index d77c5c90f..62780652a 100644
--- a/numpy/core/include/numpy/ndarraytypes.h
+++ b/numpy/core/include/numpy/ndarraytypes.h
@@ -362,10 +362,7 @@ NpyMaskValue_Create(npy_bool exposed, npy_uint8 payload)
    * allocated.
    */
 
-  /* Data buffer */
-#define PyDataMem_NEW(size) ((char *)malloc(size))
-#define PyDataMem_FREE(ptr)  free(ptr)
-#define PyDataMem_RENEW(ptr,size) ((char *)realloc(ptr,size))
+  /* Data buffer - PyDataMem_NEW/FREE/RENEW are in multiarraymodule.c */
 
 #define NPY_USE_PYMEM 1
 
@@ -1976,6 +1973,12 @@ typedef struct {
                            */
 } PyArrayInterface;
 
+/*
+ * This is a function for hooking into the PyDataMem_NEW/FREE/RENEW functions.
+ * See the documentation for PyDataMem_SetEventHook.
+ */
+typedef void (PyDataMem_EventHookFunc)(void *inp, void *outp, size_t size);
+
 #if !(defined(NPY_NO_DEPRECATED_API) && (NPY_API_VERSION <= NPY_NO_DEPRECATED_API))
 #include "npy_deprecated_api.h"
 #endif
diff --git a/numpy/core/src/multiarray/multiarray_tests.c.src b/numpy/core/src/multiarray/multiarray_tests.c.src
index 15d4c871e..09695b28f 100644
--- a/numpy/core/src/multiarray/multiarray_tests.c.src
+++ b/numpy/core/src/multiarray/multiarray_tests.c.src
@@ -375,6 +375,57 @@ clean_ax:
     return NULL;
 }
 
+/* PyDataMem_SetHook tests */
+static int malloc_count, free_count;
+static PyDataMem_EventHookFunc *old_hook = NULL;
+
+static void test_hook(void *old, void *new, size_t size)
+{
+    if (old == NULL) {
+        malloc_count++;
+    }
+    if (size == 0) {
+        free_count++;
+    }
+}
+
+static PyObject*
+test_pydatamem_seteventhook_start(PyObject* NPY_UNUSED(self), PyObject* NPY_UNUSED(args))
+{
+    malloc_count = 0;
+    free_count = 0;
+    old_hook = PyDataMem_SetEventHook(test_hook);
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
+static PyObject*
+test_pydatamem_seteventhook_end(PyObject* NPY_UNUSED(self), PyObject* NPY_UNUSED(args))
+{
+    static PyDataMem_EventHookFunc *my_hook;
+
+    my_hook = PyDataMem_SetEventHook(old_hook);
+    if (my_hook != test_hook) {
+        PyErr_SetString(PyExc_ValueError,
+                        "hook was not the expected test hook");
+        return NULL;
+    }
+
+    if (malloc_count == 0) {
+        PyErr_SetString(PyExc_ValueError,
+                        "malloc count is zero after test");
+        return NULL;
+    }
+    if (free_count == 0) {
+        PyErr_SetString(PyExc_ValueError,
+                        "free count is zero after test");
+        return NULL;
+    }
+
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
 static PyMethodDef Multiarray_TestsMethods[] = {
     {"test_neighborhood_iterator",
         test_neighborhood_iterator,
@@ -382,6 +433,12 @@ static PyMethodDef Multiarray_TestsMethods[] = {
     {"test_neighborhood_iterator_oob",
         test_neighborhood_iterator_oob,
         METH_VARARGS, NULL},
+    {"test_pydatamem_seteventhook_start",
+        test_pydatamem_seteventhook_start,
+        METH_NOARGS, NULL},
+    {"test_pydatamem_seteventhook_end",
+        test_pydatamem_seteventhook_end,
+        METH_NOARGS, NULL},
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };
 
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index 7a657d8b9..e9401240a 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -3659,6 +3659,70 @@ test_interrupt(PyObject *NPY_UNUSED(self), PyObject *args)
     return PyInt_FromLong(a);
 }
 
+/* malloc/free/realloc hook */
+NPY_NO_EXPORT PyDataMem_EventHookFunc *_PyDataMem_eventhook;
+
+/*NUMPY_API
+ * Sets the allocation event hook for numpy array data.
+ * Takes a PyDataMem_EventHookFunc *, which has the signature:
+ *        void hook(void *old, void *new, size_t size).
+ * Returns a pointer to the previous hook or NULL.
+ *
+ * If not NULL, hook will be called at the end of each PyDataMem_NEW/FREE/RENEW:
+ *   result = PyDataMem_NEW(size)        -> (*hook)(NULL, result, size)
+ *   PyDataMem_FREE(ptr)                 -> (*hook)(ptr, NULL, 0)
+ *   result = PyDataMem_RENEW(ptr, size) -> (*hook)(ptr, result, size)
+ */
+NPY_NO_EXPORT PyDataMem_EventHookFunc *
+PyDataMem_SetEventHook(PyDataMem_EventHookFunc *newhook)
+{
+    PyDataMem_EventHookFunc *temp = _PyDataMem_eventhook;
+    _PyDataMem_eventhook = newhook;
+    return temp;
+}
+
+/*NUMPY_API
+ * Allocates memory for array data.
+ */
+NPY_NO_EXPORT void *
+PyDataMem_NEW(size_t size)
+{
+    void *result;
+
+    result = malloc(size);
+    if (_PyDataMem_eventhook != NULL) {
+        (*_PyDataMem_eventhook)(NULL, result, size);
+    }
+    return (char *)result;
+}
+
+/*NUMPY_API
+ * Free memory for array data.
+ */
+NPY_NO_EXPORT void
+PyDataMem_FREE(void *ptr)
+{
+    free(ptr);
+    if (_PyDataMem_eventhook != NULL) {
+        (*_PyDataMem_eventhook)(ptr, NULL, 0);
+    }
+}
+
+/*NUMPY_API
+ * Reallocate/resize memory for array data.
+ */
+NPY_NO_EXPORT void *
+PyDataMem_RENEW(void *ptr, size_t size)
+{
+    void *result;
+
+    result = realloc(ptr, size);
+    if (_PyDataMem_eventhook != NULL) {
+        (*_PyDataMem_eventhook)(ptr, result, size);
+    }
+    return (char *)result;
+}
+
 static struct PyMethodDef array_module_methods[] = {
     {"_get_ndarray_c_version",
         (PyCFunction)array__get_ndarray_c_version,
diff --git a/numpy/core/src/npysort/sort.c.src b/numpy/core/src/npysort/sort.c.src
index 618259f96..514131748 100644
--- a/numpy/core/src/npysort/sort.c.src
+++ b/numpy/core/src/npysort/sort.c.src
@@ -38,6 +38,12 @@
 #define SMALL_MERGESORT 20
 #define SMALL_STRING 16
 
+/* These should be changed to PyDataMem_NEW/FREE if npysort is moved
+ *  to the multiarray directory.
+ */
+#define NpySortArray_malloc(size) ((char *)malloc(size))
+#define NpySortArray_free(ptr)  free(ptr)
+
 
 /*
  *****************************************************************************
@@ -335,14 +341,14 @@ mergesort_@suff@(@type@ *start, npy_intp num, void *NOT_USED)
 
     pl = start;
     pr = pl + num;
-    pw = (@type@ *) PyDataMem_NEW((num/2)*sizeof(@type@));
+    pw = (@type@ *) NpySortArray_malloc((num/2)*sizeof(@type@));
     if (!pw) {
         PyErr_NoMemory();
         return -1;
     }
     mergesort0_@suff@(pl, pr, pw);
 
-    PyDataMem_FREE(pw);
+    NpySortArray_free(pw);
     return 0;
 }
 
@@ -476,13 +482,13 @@ mergesort_@suff@(@type@ *start, npy_intp num, PyArrayObject *arr)
 
     pl = start;
     pr = pl + num*len;
-    pw = (@type@ *) PyDataMem_NEW((num/2)*elsize);
+    pw = (@type@ *) NpySortArray_malloc((num/2)*elsize);
     if (!pw) {
         PyErr_NoMemory();
         err = -1;
         goto fail_0;
     }
-    vp = (@type@ *) PyDataMem_NEW(elsize);
+    vp = (@type@ *) NpySortArray_malloc(elsize);
     if (!vp) {
         PyErr_NoMemory();
         err = -1;
@@ -490,9 +496,9 @@ mergesort_@suff@(@type@ *start, npy_intp num, PyArrayObject *arr)
     }
     mergesort0_@suff@(pl, pr, pw, vp, len);
 
-    PyDataMem_FREE(vp);
+    NpySortArray_free(vp);
 fail_1:
-    PyDataMem_FREE(pw);
+    NpySortArray_free(pw);
 fail_0:
     return err;
 }
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index e9c7a9c4f..9b975ea41 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -10,7 +10,8 @@ from numpy.testing.utils import WarningManager
 from numpy.compat import asbytes, getexception, strchar
 from test_print import in_foreign_locale
 from numpy.core.multiarray_tests import (
-        test_neighborhood_iterator, test_neighborhood_iterator_oob
+        test_neighborhood_iterator, test_neighborhood_iterator_oob,
+        test_pydatamem_seteventhook_start, test_pydatamem_seteventhook_end,
         )
 from numpy.testing import (
         TestCase, run_module_suite, assert_, assert_raises,
@@ -2633,6 +2634,16 @@ def test_flat_element_deletion():
     except:
         raise AssertionError
 
+class TestMemEventHook(TestCase):
+    def test_mem_seteventhook(self):
+        # The actual tests are within the C code in
+        # multiarray/multiarray_tests.c.src
+        test_pydatamem_seteventhook_start()
+        # force an allocation and free of a numpy array
+        a = np.zeros(10)
+        del a
+        test_pydatamem_seteventhook_end()
+
 
 if __name__ == "__main__":
     run_module_suite()
diff --git a/numpy/lib/src/_compiled_base.c b/numpy/lib/src/_compiled_base.c
index c31ee5cd8..9bb8a613d 100644
--- a/numpy/lib/src/_compiled_base.c
+++ b/numpy/lib/src/_compiled_base.c
@@ -670,7 +670,10 @@ arr_interp(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwdict)
 
     /* only pre-calculate slopes if there are relatively few of them. */
     if (lenxp <= lenx) {
-        slopes = (double *) PyDataMem_NEW((lenxp - 1)*sizeof(double));
+        slopes = (double *) PyArray_malloc((lenxp - 1)*sizeof(double));
+        if (! slopes) {
+            goto fail;
+        }
         NPY_BEGIN_ALLOW_THREADS;
         for (i = 0; i < lenxp - 1; i++) {
             slopes[i] = (dy[i + 1] - dy[i])/(dx[i + 1] - dx[i]);
@@ -692,7 +695,7 @@ arr_interp(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwdict)
             }
         }
         NPY_END_ALLOW_THREADS;
-        PyDataMem_FREE(slopes);
+        PyArray_free(slopes);
     }
     else {
         NPY_BEGIN_ALLOW_THREADS;