3 files changed, 39 insertions, 23 deletions

diff --git a/numpy/core/include/numpy/ndarraytypes.h b/numpy/core/include/numpy/ndarraytypes.h
index 62780652a..2a25479b2 100644
--- a/numpy/core/include/numpy/ndarraytypes.h
+++ b/numpy/core/include/numpy/ndarraytypes.h
@@ -1977,7 +1977,8 @@ typedef struct {
  * 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);
+typedef void (PyDataMem_EventHookFunc)(void *inp, void *outp, size_t size,
+                                       void *user_data);
 
 #if !(defined(NPY_NO_DEPRECATED_API) && (NPY_API_VERSION <= NPY_NO_DEPRECATED_API))
 #include "npy_deprecated_api.h"
diff --git a/numpy/core/src/multiarray/multiarray_tests.c.src b/numpy/core/src/multiarray/multiarray_tests.c.src
index 09695b28f..a0f70aabc 100644
--- a/numpy/core/src/multiarray/multiarray_tests.c.src
+++ b/numpy/core/src/multiarray/multiarray_tests.c.src
@@ -376,25 +376,26 @@ clean_ax:
 }
 
 /* PyDataMem_SetHook tests */
-static int malloc_count, free_count;
+static int malloc_free_counts[2];
 static PyDataMem_EventHookFunc *old_hook = NULL;
+static void *old_data;
 
-static void test_hook(void *old, void *new, size_t size)
+static void test_hook(void *old, void *new, size_t size, void *user_data)
 {
+    int* counters = (int *) user_data;
     if (old == NULL) {
-        malloc_count++;
+        counters[0]++; /* malloc counter */
     }
     if (size == 0) {
-        free_count++;
+        counters[1]++; /* free counter */
     }
 }
 
 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);
+    malloc_free_counts[0] = malloc_free_counts[1] = 0;
+    old_hook = PyDataMem_SetEventHook(test_hook, (void *) malloc_free_counts, &old_data);
     Py_INCREF(Py_None);
     return Py_None;
 }
@@ -402,21 +403,22 @@ test_pydatamem_seteventhook_start(PyObject* NPY_UNUSED(self), PyObject* NPY_UNUS
 static PyObject*
 test_pydatamem_seteventhook_end(PyObject* NPY_UNUSED(self), PyObject* NPY_UNUSED(args))
 {
-    static PyDataMem_EventHookFunc *my_hook;
+    PyDataMem_EventHookFunc *my_hook;
+    void *my_data;
 
-    my_hook = PyDataMem_SetEventHook(old_hook);
-    if (my_hook != test_hook) {
+    my_hook = PyDataMem_SetEventHook(old_hook, old_data, &my_data);
+    if ((my_hook != test_hook) || (my_data != (void *) malloc_free_counts)) {
         PyErr_SetString(PyExc_ValueError,
-                        "hook was not the expected test hook");
+                        "hook/data was not the expected test hook");
         return NULL;
     }
 
-    if (malloc_count == 0) {
+    if (malloc_free_counts[0] == 0) {
         PyErr_SetString(PyExc_ValueError,
                         "malloc count is zero after test");
         return NULL;
     }
-    if (free_count == 0) {
+    if (malloc_free_counts[1] == 0) {
         PyErr_SetString(PyExc_ValueError,
                         "free count is zero after test");
         return NULL;
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index e9401240a..3917c1611 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -3661,23 +3661,33 @@ test_interrupt(PyObject *NPY_UNUSED(self), PyObject *args)
 
 /* malloc/free/realloc hook */
 NPY_NO_EXPORT PyDataMem_EventHookFunc *_PyDataMem_eventhook;
+NPY_NO_EXPORT void *_PyDataMem_eventhook_user_data;
 
 /*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.
+ *        void hook(void *old, void *new, size_t size, void *user_data).
+ *   Also takes a void *user_data, and void **old_data.
+ *
+ * Returns a pointer to the previous hook or NULL.  If old_data is
+ * non-NULL, the previous user_data pointer will be copied to it.
+ *
  *
  * 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)
+ *   result = PyDataMem_NEW(size)        -> (*hook)(NULL, result, size, user_data)
+ *   PyDataMem_FREE(ptr)                 -> (*hook)(ptr, NULL, 0, user_data)
+ *   result = PyDataMem_RENEW(ptr, size) -> (*hook)(ptr, result, size, user_data)
  */
 NPY_NO_EXPORT PyDataMem_EventHookFunc *
-PyDataMem_SetEventHook(PyDataMem_EventHookFunc *newhook)
+PyDataMem_SetEventHook(PyDataMem_EventHookFunc *newhook,
+                       void *user_data, void **old_data)
 {
     PyDataMem_EventHookFunc *temp = _PyDataMem_eventhook;
     _PyDataMem_eventhook = newhook;
+    if (old_data != NULL) {
+        *old_data = _PyDataMem_eventhook_user_data;
+    }
+    _PyDataMem_eventhook_user_data = user_data;
     return temp;
 }
 
@@ -3691,7 +3701,8 @@ PyDataMem_NEW(size_t size)
 
     result = malloc(size);
     if (_PyDataMem_eventhook != NULL) {
-        (*_PyDataMem_eventhook)(NULL, result, size);
+        (*_PyDataMem_eventhook)(NULL, result, size,
+                                _PyDataMem_eventhook_user_data);
     }
     return (char *)result;
 }
@@ -3704,7 +3715,8 @@ PyDataMem_FREE(void *ptr)
 {
     free(ptr);
     if (_PyDataMem_eventhook != NULL) {
-        (*_PyDataMem_eventhook)(ptr, NULL, 0);
+        (*_PyDataMem_eventhook)(ptr, NULL, 0,
+                                _PyDataMem_eventhook_user_data);
     }
 }
 
@@ -3718,7 +3730,8 @@ PyDataMem_RENEW(void *ptr, size_t size)
 
     result = realloc(ptr, size);
     if (_PyDataMem_eventhook != NULL) {
-        (*_PyDataMem_eventhook)(ptr, result, size);
+        (*_PyDataMem_eventhook)(ptr, result, size,
+                                _PyDataMem_eventhook_user_data);
     }
     return (char *)result;
 }