forked from numpy/numpy
/
alloc_hook.pyx
42 lines (34 loc) · 1.49 KB
/
alloc_hook.pyx
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
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)