8000 Merge pull request #24248 from seberg/getenv-slow-warn-mem-policy · numpy/numpy@f1eb09d · GitHub
[go: up one dir, main page]

Skip to content

Commit f1eb09d

Browse files
authored
Merge pull request #24248 from seberg/getenv-slow-warn-mem-policy
BUG: Factor out slow `getenv` call used for memory policy warning
2 parents 1c51f77 + 1d0487b commit f1eb09d

File tree

5 files changed

+49
-13
lines changed

5 files changed

+49
-13
lines changed

doc/source/reference/global_state.rst

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,3 +82,6 @@ memory allocation policy, the default will be to call ``free``. If
8282
``NUMPY_WARN_IF_NO_MEM_POLICY`` is set to ``"1"``, a ``RuntimeWarning`` will
8383
be emitted. A better alternative is to use a ``PyCapsule`` with a deallocator
8484
and set the ``ndarray.base``.
85+
86+
.. versionchanged:: 1.25.2
87+
This variable is only checked on the first import.

numpy/core/src/multiarray/arrayobject.c

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,9 @@ maintainer email: oliphant.travis@ieee.org
6262
#include "binop_override.h"
6363
#include "array_coercion.h"
6464

65+
66+
NPY_NO_EXPORT npy_bool numpy_warn_if_no_mem_policy = 0;
67+
6568
/*NUMPY_API
6669
Compute the size of an array (in number of items)
6770
*/
@@ -460,9 +463,8 @@ array_dealloc(PyArrayObject *self)
460463
}
461464
}
462465
if (fa->mem_handler == NULL) {
463-
char *env = getenv("NUMPY_WARN_IF_NO_MEM_POLICY");
464-
if ((env != NULL) && (strncmp(env, "1", 1) == 0)) {
465-
char const * msg = "Trying to dealloc data, but a memory policy "
466+
if (numpy_warn_if_no_mem_policy) {
467+
char const *msg = "Trying to dealloc data, but a memory policy "
466468
"is not set. If you take ownership of the data, you must "
467469
"set a base owning the data (e.g. a PyCapsule).";
468470
WARN_IN_DEALLOC(PyExc_RuntimeWarning, msg);

numpy/core/src/multiarray/arrayobject.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
#ifndef NUMPY_CORE_SRC_MULTIARRAY_ARRAYOBJECT_H_
66
#define NUMPY_CORE_SRC_MULTIARRAY_ARRAYOBJECT_H_
77

8+
extern NPY_NO_EXPORT npy_bool numpy_warn_if_no_mem_policy;
9+
810
NPY_NO_EXPORT PyObject *
911
_strings_richcompare(PyArrayObject *self, PyArrayObject *other, int cmp_op,
1012
int rstrip);

numpy/core/src/multiarray/multiarraymodule.c

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4378,6 +4378,24 @@ normalize_axis_index(PyObject *NPY_UNUSED(self),
43784378
}
43794379

43804380

4381+
static PyObject *
4382+
_set_numpy_warn_if_no_mem_policy(PyObject *NPY_UNUSED(self), PyObject *arg)
4383+
{
4384+
int res = PyObject_IsTrue(arg);
4385+
if (res < 0) {
4386+
return NULL;
4387+
}
4388+
int old_value = numpy_warn_if_no_mem_policy;
4389+
numpy_warn_if_no_mem_policy = res;
4390+
if (old_value) {
4391+
Py_RETURN_TRUE;
4392+
}
4393+
else {
4394+
Py_RETURN_FALSE;
4395+
}
4396+
}
4397+
4398+
43814399
static PyObject *
43824400
_reload_guard(PyObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args)) {
43834401
static int initialized = 0;
@@ -4625,6 +4643,9 @@ static struct PyMethodDef array_module_methods[] = {
46254643
METH_O, "Set the NEP 50 promotion state. This is not thread-safe.\n"
46264644
"The optional warnings can be safely silenced using the \n"
46274645
"`np._no_nep50_warning()` context manager."},
4646+
{"_set_numpy_warn_if_no_mem_policy",
4647+
(PyCFunction)_set_numpy_warn_if_no_mem_policy,
4648+
METH_O, "Change the warn if no mem policy flag for testing."},
46284649
{"_add_newdoc_ufunc", (PyCFunction)add_newdoc_ufunc,
46294650
METH_VARARGS, NULL},
46304651
{"_get_sfloat_dtype",
@@ -4913,6 +4934,14 @@ initialize_static_globals(void)
49134934
return -1;
49144935
}
49154936

4937+
char *env = getenv("NUMPY_WARN_IF_NO_MEM_POLICY");
4938+
if ((env != NULL) && (strncmp(env, "1", 1) == 0)) {
4939+
numpy_warn_if_no_mem_policy = 1;
4940+
}
4941+
else {
4942+
numpy_warn_if_no_mem_policy = 0;
4943+
}
4944+
49164945
return 0;
49174946
}
49184947

numpy/core/tests/test_mem_policy.py

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -409,16 +409,19 @@ def test_switch_owner(get_module, policy):
409409
a = get_module.get_array()
410410
assert np.core.multiarray.get_handler_name(a) is None
411411
get_module.set_own(a)
412-
oldval = os.environ.get('NUMPY_WARN_IF_NO_MEM_POLICY', None)
412+
413413
if policy is None:
414-
if 'NUMPY_WARN_IF_NO_MEM_POLICY' in os.environ:
415-
os.environ.pop('NUMPY_WARN_IF_NO_MEM_POLICY')
414+
# See what we expect to be set based on the env variable
415+
policy = os.getenv("NUMPY_WARN_IF_NO_MEM_POLICY", "0") == "1"
416+
oldval = None
416417
else:
417-
os.environ['NUMPY_WARN_IF_NO_MEM_POLICY'] = policy
418+
policy = policy == "1"
419+
oldval = np.core._multiarray_umath._set_numpy_warn_if_no_mem_policy(
420+
policy)
418421
try:
419422
# The policy should be NULL, so we have to assume we can call
420423
# "free". A warning is given if the policy == "1"
421-
if policy == "1":
424+
if policy:
422425
with assert_warns(RuntimeWarning) as w:
423426
del a
424427
gc.collect()
@@ -427,11 +430,8 @@ def test_switch_owner(get_module, policy):
427430
gc.collect()
428431

429432
finally:
430-
if oldval is None:
431-
if 'NUMPY_WARN_IF_NO_MEM_POLICY' in os.environ:
432-
os.environ.pop('NUMPY_WARN_IF_NO_MEM_POLICY')
433-
else:
434-
os.environ['NUMPY_WARN_IF_NO_MEM_POLICY'] = oldval
433+
if oldval is not None:
434+
np.core._multiarray_umath._set_numpy_warn_if_no_mem_policy(oldval)
435435

436436

437437
@pytest.mark.skipif(sys.version_info >= (3, 12), reason="no numpy.distutils")

0 commit comments

Comments
 (0)
0