diff --git a/doc/source/reference/c-api/data_memory.rst b/doc/source/reference/c-api/data_memory.rst index 2084ab5d083f..5cf439912699 100644 --- a/doc/source/reference/c-api/data_memory.rst +++ b/doc/source/reference/c-api/data_memory.rst @@ -55,6 +55,8 @@ functions may change during the lifetime of the process, each ``ndarray`` carries with it the functions used at the time of its instantiation, and these will be used to reallocate or free the data memory of the instance. +For details see: :ref:`NEP 49 — Data allocation strategies `. + .. c:type:: PyDataMem_Handler A struct to hold function pointers used to manipulate memory @@ -87,6 +89,9 @@ will be used to reallocate or free the data memory of the instance. return ``NULL`` if an error has occurred. We wrap the user-provided functions so they will still call the python and numpy memory management callback hooks. + The handlers are stored in a Python context variable + (see https://docs.python.org/3/library/contextvars.html), + so there can be multiple handlers in a Python session. .. c:function:: PyObject * PyDataMem_GetHandler() diff --git a/numpy/core/src/multiarray/alloc.c b/numpy/core/src/multiarray/alloc.c index 8b765aa954da..a19e9196ff8b 100644 --- a/numpy/core/src/multiarray/alloc.c +++ b/numpy/core/src/multiarray/alloc.c @@ -430,11 +430,23 @@ int uo_index=0; /* user_override index */ /* Wrappers for the default or any user-assigned PyDataMem_Handler */ +static inline PyDataMem_Handler * +_PyDataMem_GetHandler_Internal(PyObject *mem_handler) +{ + if (mem_handler == PyDataMem_DefaultHandler) + // fast path for default allocator + return &default_handler; + + PyDataMem_Handler *handler = (PyDataMem_Handler *)PyCapsule_GetPointer( + mem_handler, "mem_handler"); + return handler; +} + NPY_NO_EXPORT void * PyDataMem_UserNEW(size_t size, PyObject *mem_handler) { void *result; - PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler"); + PyDataMem_Handler *handler = _PyDataMem_GetHandler_Internal(mem_handler); if (handler == NULL) { return NULL; } @@ -457,7 +469,7 @@ NPY_NO_EXPORT void * PyDataMem_UserNEW_ZEROED(size_t nmemb, size_t size, PyObject *mem_handler) { void *result; - PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler"); + PyDataMem_Handler *handler = _PyDataMem_GetHandler_Internal(mem_handler); if (handler == NULL) { return NULL; } @@ -479,7 +491,7 @@ PyDataMem_UserNEW_ZEROED(size_t nmemb, size_t size, PyObject *mem_handler) NPY_NO_EXPORT void PyDataMem_UserFREE(void *ptr, size_t size, PyObject *mem_handler) { - PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler"); + PyDataMem_Handler *handler = _PyDataMem_GetHandler_Internal(mem_handler); if (handler == NULL) { WARN_NO_RETURN(PyExc_RuntimeWarning, "Could not get pointer to 'mem_handler' from PyCapsule"); @@ -502,7 +514,7 @@ NPY_NO_EXPORT void * PyDataMem_UserRENEW(void *ptr, size_t size, PyObject *mem_handler) { void *result; - PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler"); + PyDataMem_Handler *handler = _PyDataMem_GetHandler_Internal(mem_handler); if (handler == NULL) { return NULL; }