summaryrefslogtreecommitdiff
path: root/numpy/core
diff options
context:
space:
mode:
Diffstat (limited to 'numpy/core')
-rw-r--r--numpy/core/_add_newdocs.py10
-rw-r--r--numpy/core/include/numpy/ndarraytypes.h7
-rw-r--r--numpy/core/multiarray.py16
-rw-r--r--numpy/core/src/multiarray/alloc.c39
-rw-r--r--numpy/core/src/multiarray/alloc.h2
-rw-r--r--numpy/core/src/multiarray/multiarraymodule.c3
-rw-r--r--numpy/core/tests/test_mem_policy.py6
7 files changed, 73 insertions, 10 deletions
diff --git a/numpy/core/_add_newdocs.py b/numpy/core/_add_newdocs.py
index cae5bc281..078c58976 100644
--- a/numpy/core/_add_newdocs.py
+++ b/numpy/core/_add_newdocs.py
@@ -4759,6 +4759,16 @@ add_newdoc('numpy.core.multiarray', 'get_handler_name',
memory, in which case you can traverse ``a.base`` for a memory handler.
""")
+add_newdoc('numpy.core.multiarray', 'get_handler_version',
+ """
+ get_handler_version(a: ndarray) -> int,None
+
+ Return the version of the memory handler used by `a`. If not provided,
+ return the version 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/include/numpy/ndarraytypes.h b/numpy/core/include/numpy/ndarraytypes.h
index 80177e2bb..2607fb732 100644
--- a/numpy/core/include/numpy/ndarraytypes.h
+++ b/numpy/core/include/numpy/ndarraytypes.h
@@ -674,10 +674,15 @@ typedef struct {
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);
+ /*
+ * This is the end of the version=1 struct. Only add new fields after
+ * this line
+ */
} PyDataMemAllocator;
typedef struct {
- char name[128]; /* multiple of 64 to keep the struct aligned */
+ char name[127]; /* multiple of 64 to keep the struct aligned */
+ uint8_t version; /* currently 1 */
PyDataMemAllocator allocator;
} PyDataMem_Handler;
diff --git a/numpy/core/multiarray.py b/numpy/core/multiarray.py
index f96274263..f88d75978 100644
--- a/numpy/core/multiarray.py
+++ b/numpy/core/multiarray.py
@@ -33,14 +33,14 @@ __all__ = [
'dot', 'dragon4_positional', 'dragon4_scientific', 'dtype',
'empty', 'empty_like', 'error', 'flagsobj', 'flatiter', 'format_longfloat',
'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',
- 'set_datetimeparse_function', 'set_legacy_print_mode', 'set_numeric_ops',
- 'set_string_function', 'set_typeDict', 'shares_memory',
- 'tracemalloc_domain', 'typeinfo', 'unpackbits', 'unravel_index', 'vdot',
- 'where', 'zeros']
+ 'get_handler_name', 'get_handler_version', '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', 'set_datetimeparse_function',
+ 'set_legacy_print_mode', 'set_numeric_ops', 'set_string_function',
+ 'set_typeDict', 'shares_memory', 'tracemalloc_domain', 'typeinfo',
+ 'unpackbits', 'unravel_index', 'vdot', 'where', 'zeros']
# For backward compatibility, make sure pickle imports these functions from here
_reconstruct.__module__ = 'numpy.core.multiarray'
diff --git a/numpy/core/src/multiarray/alloc.c b/numpy/core/src/multiarray/alloc.c
index e4756264d..d1173410d 100644
--- a/numpy/core/src/multiarray/alloc.c
+++ b/numpy/core/src/multiarray/alloc.c
@@ -370,6 +370,7 @@ default_free(void *NPY_UNUSED(ctx), void *ptr, size_t size)
/* Memory handler global default */
PyDataMem_Handler default_handler = {
"default_allocator",
+ 1,
{
NULL, /* ctx */
default_malloc, /* malloc */
@@ -395,7 +396,6 @@ PyDataMem_UserNEW(size_t size, PyObject *mem_handler)
if (handler == NULL) {
return NULL;
}
-
assert(size != 0);
result = handler->allocator.malloc(handler->allocator.ctx, size);
if (_PyDataMem_eventhook != NULL) {
@@ -639,3 +639,40 @@ get_handler_name(PyObject *NPY_UNUSED(self), PyObject *args)
Py_DECREF(mem_handler);
return name;
}
+
+NPY_NO_EXPORT PyObject *
+get_handler_version(PyObject *NPY_UNUSED(self), PyObject *args)
+{
+ PyObject *arr=NULL;
+ if (!PyArg_ParseTuple(args, "|O:get_handler_version", &arr)) {
+ return NULL;
+ }
+ if (arr != NULL && !PyArray_Check(arr)) {
+ PyErr_SetString(PyExc_ValueError, "if supplied, argument must be an ndarray");
+ return NULL;
+ }
+ PyObject *mem_handler;
+ PyDataMem_Handler *handler;
+ PyObject *version;
+ if (arr != NULL) {
+ mem_handler = PyArray_HANDLER((PyArrayObject *) arr);
+ if (mem_handler == NULL) {
+ Py_RETURN_NONE;
+ }
+ Py_INCREF(mem_handler);
+ }
+ else {
+ mem_handler = PyDataMem_GetHandler();
+ if (mem_handler == NULL) {
+ return NULL;
+ }
+ }
+ handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+ if (handler == NULL) {
+ Py_DECREF(mem_handler);
+ return NULL;
+ }
+ version = PyLong_FromLong(handler->version);
+ Py_DECREF(mem_handler);
+ return version;
+}
diff --git a/numpy/core/src/multiarray/alloc.h b/numpy/core/src/multiarray/alloc.h
index 4f7df1f84..f1ccf0bcd 100644
--- a/numpy/core/src/multiarray/alloc.h
+++ b/numpy/core/src/multiarray/alloc.h
@@ -47,5 +47,7 @@ extern PyDataMem_Handler default_handler;
NPY_NO_EXPORT PyObject *
get_handler_name(PyObject *NPY_UNUSED(self), PyObject *obj);
+NPY_NO_EXPORT PyObject *
+get_handler_version(PyObject *NPY_UNUSED(self), PyObject *obj);
#endif /* NUMPY_CORE_SRC_MULTIARRAY_ALLOC_H_ */
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index 84179d5f0..a854bcb3b 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -4437,6 +4437,9 @@ static struct PyMethodDef array_module_methods[] = {
{"get_handler_name",
(PyCFunction) get_handler_name,
METH_VARARGS, NULL},
+ {"get_handler_version",
+ (PyCFunction) get_handler_version,
+ METH_VARARGS, NULL},
{"_add_newdoc_ufunc", (PyCFunction)add_newdoc_ufunc,
METH_VARARGS, NULL},
{"_get_sfloat_dtype",
diff --git a/numpy/core/tests/test_mem_policy.py b/numpy/core/tests/test_mem_policy.py
index 7fec8897f..abf340062 100644
--- a/numpy/core/tests/test_mem_policy.py
+++ b/numpy/core/tests/test_mem_policy.py
@@ -179,6 +179,7 @@ def get_module(tmp_path):
};
static PyDataMem_Handler secret_data_handler = {
"secret_data_allocator",
+ 1,
{
&secret_data_handler_ctx, /* ctx */
shift_alloc, /* malloc */
@@ -212,17 +213,22 @@ def get_module(tmp_path):
def test_set_policy(get_module):
get_handler_name = np.core.multiarray.get_handler_name
+ get_handler_version = np.core.multiarray.get_handler_version
orig_policy_name = get_handler_name()
a = np.arange(10).reshape((2, 5)) # a doesn't own its own data
assert get_handler_name(a) is None
+ assert get_handler_version(a) is None
assert get_handler_name(a.base) == orig_policy_name
+ assert get_handler_version(a.base) == 1
orig_policy = get_module.set_secret_data_policy()
b = np.arange(10).reshape((2, 5)) # b doesn't own its own data
assert get_handler_name(b) is None
+ assert get_handler_version(b) is None
assert get_handler_name(b.base) == 'secret_data_allocator'
+ assert get_handler_version(b.base) == 1
if orig_policy_name == 'default_allocator':
get_module.set_old_policy(None) # tests PyDataMem_SetHandler(NULL)
View Lorry Controller Status