Wrap hook functions with GIL, add example.

Wraps the SetHook and calls to the hook with the GIL, to prevent races.
Adds an example of using the interface for callbacks into python code.

1 files changed, 42 insertions, 0 deletions

diff --git a/tools/allocation_tracking/alloc_hook.pyx b/tools/allocation_tracking/alloc_hook.pyx
new file mode 100644
index 000000000..d1e656f90
--- /dev/null
+++ b/tools/allocation_tracking/alloc_hook.pyx
@@ -0,0 +1,42 @@
+# A cython wrapper for using python functions as callbacks for
+# PyDataMem_SetEventHook.
+
+cimport numpy as np
+
+cdef extern from "Python.h":
+    object PyLong_FromVoidPtr(void *)
+    void *PyLong_AsVoidPtr(object)
+
+ctypedef void PyDataMem_EventHookFunc(void *inp, void *outp, size_t size,
+                                      void *user_data)
+cdef extern from "numpy/arrayobject.h":
+    PyDataMem_EventHookFunc * \
+        PyDataMem_SetEventHook(PyDataMem_EventHookFunc *newhook,
+                               void *user_data, void **old_data)
+
+np.import_array()
+
+cdef void pyhook(void *old, void *new, size_t size, void *user_data):
+    cdef object pyfunc = <object> user_data
+    pyfunc(PyLong_FromVoidPtr(old),
+           PyLong_FromVoidPtr(new),
+           size)
+
+class NumpyAllocHook(object):
+    def __init__(self, callback):
+        self.callback = callback
+
+    def __enter__(self):
+        cdef void *old_hook, *old_data
+        old_hook = <void *> \
+            PyDataMem_SetEventHook(<PyDataMem_EventHookFunc *> pyhook,
+                                    <void *> self.callback,
+                                    <void **> &old_data)
+        self.old_hook = PyLong_FromVoidPtr(old_hook)
+        self.old_data = PyLong_FromVoidPtr(old_data)
+
+    def __exit__(self):
+        PyDataMem_SetEventHook(<PyDataMem_EventHookFunc *> \
+                                    PyLong_AsVoidPtr(self.old_hook),
+                                <void *> PyLong_AsVoidPtr(self.old_data),
+                                <void **> 0)