numpy · seberg · Oct 25, 2021 · Oct 11, 2020 · Oct 16, 2020 · Oct 18, 2020
diff --git a/doc/TESTS.rst.txt b/doc/TESTS.rst.txt
@@ -139,6 +139,21 @@ originally written without unit tests, there are still several modules
 that don't have tests yet. Please feel free to choose one of these
 modules and develop tests for it.
 
+Using C code in tests
+---------------------
+
+NumPy exposes a rich :ref:`C-API<c-api>` . These are tested using c-extension
+modules written "as-if" they know nothing about the internals of NumPy, rather
+using the official C-API interfaces only. Examples of such modules are tests
+for a user-defined ``rational`` dtype in ``_rational_tests`` or the ufunc
+machinery tests in ``_umath_tests`` which are part of the binary distribution.
+Starting from version 1.21, you can also write snippets of C code in tests that
+will be compiled locally into c-extension modules and loaded into python.
+
+.. currentmodule:: numpy.testing.extbuild
+
+.. autofunction:: build_and_import_extension
+
 Labeling tests
 --------------
 

diff --git a/doc/neps/nep-0049.rst b/doc/neps/nep-0049.rst
@@ -93,19 +93,21 @@ High level design
 
 Users who wish to change the NumPy data memory management routines will use
 :c:func:`PyDataMem_SetHandler`, which uses a :c:type:`PyDataMem_Handler`
-structure to hold pointers to functions used to manage the data memory.
+structure to hold pointers to functions used to manage the data memory. In
+order to allow lifetime management of the ``context``, the structure is wrapped
+in a ``PyCapsule``.
 
 Since a call to ``PyDataMem_SetHandler`` will change the default functions, but
 that function may be called during the lifetime of an ``ndarray`` object, each
-``ndarray`` will carry with it the ``PyDataMem_Handler`` struct used at the
-time of its instantiation, and these will be used to reallocate or free the
-data memory of the instance. Internally NumPy may use ``memcpy`` or ``memset``
-on the pointer to the data memory.
+``ndarray`` will carry with it the ``PyDataMem_Handler``-wrapped PyCapsule used
+at the time of its instantiation, and these will be used to reallocate or free
+the data memory of the instance. Internally NumPy may use ``memcpy`` or
+``memset`` on the pointer to the data memory.
 
 The name of the handler will be exposed on the python level via a
 ``numpy.core.multiarray.get_handler_name(arr)`` function. If called as
 ``numpy.core.multiarray.get_handler_name()`` it will return the name of the
-global handler that will be used to allocate data for the next new `ndarrray`.
+handler that will be used to allocate data for the next new `ndarrray`.
 
 NumPy C-API functions
 =====================
@@ -150,20 +152,19 @@ NumPy C-API functions
     15780_ and 15788_ but has not yet been resolved. When it is this NEP should
     be revisited.
 
-.. c:function:: const PyDataMem_Handler * PyDataMem_SetHandler(PyDataMem_Handler *handler)
+.. c:function:: PyObject * PyDataMem_SetHandler(PyObject *handler)
 
    Sets a new allocation policy. If the input value is ``NULL``, will reset
-   the policy to the default. Returns the previous policy, ``NULL`` if the
-   previous policy was the default. We wrap the user-provided functions
+   the policy to the default. Return the previous policy, or
+   return NULL if an error has occurred. We wrap the user-provided
    so they will still call the Python and NumPy memory management callback
    hooks. All the function pointers must be filled in, ``NULL`` is not
    accepted.
 
-.. c:function:: const PyDataMem_Handler * PyDataMem_GetHandler(PyArrayObject *obj)
+.. c:function:: const PyObject * PyDataMem_GetHandler()
 
-   Return the ``PyDataMem_Handler`` used by the
-   ``PyArrayObject``. If ``NULL``, return the handler
-   that will be used to allocate data for the next ``PyArrayObject``.
+   Return the current policy that will be used to allocate data for the
+   next ``PyArrayObject``. On failure, return ``NULL``.
 
 ``PyDataMem_Handler`` thread safety and lifetime
 ================================================

diff --git a/doc/source/reference/c-api/data_memory.rst b/doc/source/reference/c-api/data_memory.rst
@@ -0,0 +1,119 @@
+Memory management in NumPy
+==========================
+
+The `numpy.ndarray` is a python class. It requires additional memory allocations
+to hold `numpy.ndarray.strides`, `numpy.ndarray.shape` and
+`numpy.ndarray.data` attributes. These attributes are specially allocated
+after creating the python object in `__new__`. The ``strides`` and
+``shape`` are stored in a piece of memory allocated internally.
+
+The ``data`` allocation used to store the actual array values (which could be
+pointers in the case of ``object`` arrays) can be very large, so NumPy has
+provided interfaces to manage its allocation and release. This document details
+how those interfaces work.
+
+Historical overview
+-------------------
+
+Since version 1.7.0, NumPy has exposed a set of ``PyDataMem_*`` functions
+(:c:func:`PyDataMem_NEW`, :c:func:`PyDataMem_FREE`, :c:func:`PyDataMem_RENEW`)
+which are backed by `alloc`, `free`, `realloc` respectively. In that version
+NumPy also exposed the `PyDataMem_EventHook` function described below, which
+wrap the OS-level calls.
+
+Since those early days, Python also improved its memory management
+capabilities, and began providing
+various :ref:`management policies <memoryoverview>` beginning in version
+3.4. These routines are divided into a set of domains, each domain has a
+:c:type:`PyMemAllocatorEx` structure of routines for memory management. Python also
+added a `tracemalloc` module to trace calls to the various routines. These
+tracking hooks were added to the NumPy ``PyDataMem_*`` routines.
+
+NumPy added a small cache of allocated memory in its internal
+``npy_alloc_cache``, ``npy_alloc_cache_zero``, and ``npy_free_cache``
+functions. These wrap ``alloc``, ``alloc-and-memset(0)`` and ``free``
+respectively, but when ``npy_free_cache`` is called, it adds the pointer to a
+short list of available blocks marked by size. These blocks can be re-used by
+subsequent calls to ``npy_alloc*``, avoiding memory thrashing.
+
+Configurable memory routines in NumPy (NEP 49)
+----------------------------------------------
+
+Users may wish to override the internal data memory routines with ones of their
+own. Since NumPy does not use the Python domain strategy to manage data memory,
+it provides an alternative set of C-APIs to change memory routines. There are
+no Python domain-wide strategies for large chunks of object data, so those are
+less suited to NumPy's needs. User who wish to change the NumPy data memory
+management routines can use :c:func:`PyDataMem_SetHandler`, which uses a
+:c:type:`PyDataMem_Handler` structure to hold pointers to functions used to
+manage the data memory. The calls are still wrapped by internal routines to
+call :c:func:`PyTraceMalloc_Track`, :c:func:`PyTraceMalloc_Untrack`, and will
+use the :c:func:`PyDataMem_EventHookFunc` mechanism. Since the 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.
+
+.. c:type:: PyDataMem_Handler
+
+    A struct to hold function pointers used to manipulate memory
+
+    .. code-block:: c
+
+        typedef struct {
+            char name[128];  /* multiple of 64 to keep the struct aligned */
+            PyDataMemAllocator allocator;
+        } PyDataMem_Handler;
+
+    where the allocator structure is
+
+    .. code-block:: c
+
+        /* The declaration of free differs from PyMemAllocatorEx */ 
+        typedef struct {
+            void *ctx;
+            void* (*malloc) (void *ctx, size_t size);
+            void* (*calloc) (void *ctx, size_t nelem, size_t elsize);
+            void* (*realloc) (void *ctx, void *ptr, size_t new_size);
+            void (*free) (void *ctx, void *ptr, size_t size);
+        } PyDataMemAllocator;
+
+.. c:function:: PyObject * PyDataMem_SetHandler(PyObject *handler)
+
+   Set a new allocation policy. If the input value is ``NULL``, will reset the
+   policy to the default. Return the previous policy, or
+   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.
+
+.. c:function:: PyObject * PyDataMem_GetHandler()
+
+   Return the current policy that will be used to allocate data for the
+   next ``PyArrayObject``. On failure, return ``NULL``.
+
+For an example of setting up and using the PyDataMem_Handler, see the test in
+:file:`numpy/core/tests/test_mem_policy.py`
+
+.. c:function:: void PyDataMem_EventHookFunc(void *inp, void *outp, size_t size, void *user_data);
+
+    This function will be called during data memory manipulation
+
+.. c:function:: PyDataMem_EventHookFunc * PyDataMem_SetEventHook(PyDataMem_EventHookFunc *newhook, void *user_data, void **old_data)
+
+    Sets the allocation event hook for numpy array 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``:
+
+    .. code-block:: c
+
+        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)
+
+    When the hook is called, the GIL will be held by the calling
+    thread.  The hook should be written to be reentrant, if it performs
+    operations that might cause new allocation events (such as the
+    creation/destruction numpy objects, or creating/destroying Python
+    objects which might cause a gc)
diff --git a/doc/source/reference/c-api/index.rst b/doc/source/reference/c-api/index.rst
@@ -49,3 +49,4 @@ code.
    generalized-ufuncs
    coremath
    deprecations
+   data_memory
diff --git a/numpy/core/_add_newdocs.py b/numpy/core/_add_newdocs.py
@@ -4727,6 +4727,16 @@
     and then throwing away the ufunc.
     """)
 
+add_newdoc('numpy.core.multiarray', 'get_handler_name',
+    """
+    get_handler_name(a: ndarray) -> str,None
+
+    Return the name of the memory handler used by `a`. If not provided, return
+    the name of the memory handler that will be used to allocate data for the
+    next `ndarray` in this context. May return None if `a` does not own its
+    memory, in which case you can traverse ``a.base`` for a memory handler.
+    """)
+
 add_newdoc('numpy.core.multiarray', '_set_madvise_hugepage',
     """
     _set_madvise_hugepage(enabled: bool) -> bool

diff --git a/numpy/core/code_generators/cversions.txt b/numpy/core/code_generators/cversions.txt
@@ -56,5 +56,7 @@
 # DType related API additions.
 # A new field was added to the end of PyArrayObject_fields.
 # Version 14 (NumPy 1.21) No change.
-# Version 14 (NumPy 1.22) No change.
 0x0000000e = 17a0f366e55ec05e5c5c149123478452
+
+# Version 15 (NumPy 1.22) Configurable memory allocations
+0x0000000f = 0c420aed67010594eb81f23ddfb02a88
diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
@@ -76,9 +76,9 @@
     # End 1.6 API
 }
 
-#define NPY_NUMUSERTYPES (*(int *)PyArray_API[6])
-#define PyBoolArrType_Type (*(PyTypeObject *)PyArray_API[7])
-#define _PyArrayScalar_BoolValues ((PyBoolScalarObject *)PyArray_API[8])
+# define NPY_NUMUSERTYPES (*(int *)PyArray_API[6])
+# define PyBoolArrType_Type (*(PyTypeObject *)PyArray_API[7])
+# define _PyArrayScalar_BoolValues ((PyBoolScalarObject *)PyArray_API[8])
 
 multiarray_funcs_api = {
     'PyArray_GetNDArrayCVersion':           (0,),
@@ -350,6 +350,9 @@
     'PyArray_ResolveWritebackIfCopy':       (302,),
     'PyArray_SetWritebackIfCopyBase':       (303,),
     # End 1.14 API
+    'PyDataMem_SetHandler':                 (304,),
+    'PyDataMem_GetHandler':                 (305,),
+    # End 1.21 API
 }
 
 ufunc_types_api = {

diff --git a/numpy/core/include/numpy/ndarraytypes.h b/numpy/core/include/numpy/ndarraytypes.h
@@ -355,12 +355,10 @@ struct NpyAuxData_tag {
 #define NPY_ERR(str) fprintf(stderr, #str); fflush(stderr);
 #define NPY_ERR2(str) fprintf(stderr, str); fflush(stderr);
 
-  /*
-   * Macros to define how array, and dimension/strides data is
-   * allocated.
-   */
-
-  /* Data buffer - PyDataMem_NEW/FREE/RENEW are in multiarraymodule.c */
+/*
+* Macros to define how array, and dimension/strides data is
+* allocated. These should be made private
+*/
 
 #define NPY_USE_PYMEM 1
 
@@ -666,6 +664,24 @@ typedef struct _arr_descr {
         PyObject *shape;       /* a tuple */
 } PyArray_ArrayDescr;
 
+/*
+ * Memory handler structure for array data.
+ */
+/* The declaration of free differs from PyMemAllocatorEx */
+typedef struct {
+    void *ctx;
+    void* (*malloc) (void *ctx, size_t size);
+    void* (*calloc) (void *ctx, size_t nelem, size_t elsize);
+    void* (*realloc) (void *ctx, void *ptr, size_t new_size);
+    void (*free) (void *ctx, void *ptr, size_t size);
+} PyDataMemAllocator;
+
+typedef struct {
+    char name[128];  /* multiple of 64 to keep the struct aligned */
+    PyDataMemAllocator allocator;
+} PyDataMem_Handler;
+
+
 /*
  * The main array object structure.
  *
@@ -716,6 +732,10 @@ typedef struct tagPyArrayObject_fields {
     /* For weak references */
     PyObject *weakreflist;
     void *_buffer_info;  /* private buffer info, tagged to allow warning */
+    /*
+     * For malloc/calloc/realloc/free per object
+     */
+    PyObject *mem_handler;
 } PyArrayObject_fields;
 
 /*
@@ -1659,6 +1679,12 @@ PyArray_CLEARFLAGS(PyArrayObject *arr, int flags)
     ((PyArrayObject_fields *)arr)->flags &= ~flags;
 }
 
+static NPY_INLINE NPY_RETURNS_BORROWED_REF PyObject *
+PyArray_HANDLER(PyArrayObject *arr)
+{
+    return ((PyArrayObject_fields *)arr)->mem_handler;
+}
+
 #define PyTypeNum_ISBOOL(type) ((type) == NPY_BOOL)
 
 #define PyTypeNum_ISUNSIGNED(type) (((type) == NPY_UBYTE) ||   \

diff --git a/numpy/core/multiarray.py b/numpy/core/multiarray.py
@@ -31,8 +31,8 @@
     'count_nonzero', 'c_einsum', 'datetime_as_string', 'datetime_data',
     'dot', 'dragon4_positional', 'dragon4_scientific', 'dtype',
     'empty', 'empty_like', 'error', 'flagsobj', 'flatiter', 'format_longfloat',
-    'frombuffer', 'fromfile', 'fromiter', 'fromstring', 'inner',
-    'interp', 'interp_complex', 'is_busday', 'lexsort',
+    'frombuffer', 'fromfile', 'fromiter', 'fromstring', 'get_handler_name',
-    'frombuffer', 'fromfile', 'fromiter', 'fromstring', 'get_handler_name',
+    'frombuffer', 'fromfile', 'fromiter', 'fromstring', 'get_handler_name',
-    'frombuffer', 'fromfile', 'fromiter', 'fromstring', 'get_handler_name',
+    'frombuffer', 'fromfile', 'fromiter', 'fromstring', 'get_handler_name',
+    'inner', 'interp', 'interp_complex', 'is_busday', 'lexsort',
     'matmul', 'may_share_memory', 'min_scalar_type', 'ndarray', 'nditer',
     'nested_iters', 'normalize_axis_index', 'packbits',
     'promote_types', 'putmask', 'ravel_multi_index', 'result_type', 'scalar',

diff --git a/numpy/core/setup_common.py b/numpy/core/setup_common.py
@@ -43,8 +43,8 @@
 # 0x0000000d - 1.19.x
 # 0x0000000e - 1.20.x
 # 0x0000000e - 1.21.x
-# 0x0000000e - 1.22.x
-C_API_VERSION = 0x0000000e
+# 0x0000000f - 1.22.x
+C_API_VERSION = 0x0000000f
 
 class MismatchCAPIWarning(Warning):
     pass