File: alloc_hook.pyx

package info (click to toggle)
numpy 1%3A1.19.5-1
  • links: PTS, VCS
  • area: main
  • in suites: bullseye
  • size: 27,552 kB
  • sloc: ansic: 164,908; python: 128,463; cpp: 1,117; makefile: 594; javascript: 387; f90: 298; sh: 294; fortran: 200; sed: 140; perl: 34
file content (42 lines) | stat: -rw-r--r-- 1,521 bytes parent folder | download
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:
    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)