numpy · mattip · Jul 24, 2023 · Jul 24, 2023
diff --git a/doc/source/reference/global_state.rst b/doc/source/reference/global_state.rst
@@ -82,3 +82,6 @@ memory allocation policy, the default will be to call ``free``. If
 ``NUMPY_WARN_IF_NO_MEM_POLICY`` is set to ``"1"``, a ``RuntimeWarning`` will
 be emitted. A better alternative is to use a ``PyCapsule`` with a deallocator
 and set the ``ndarray.base``.
+
+.. versionchanged:: 1.25.2
+    This variable is only checked on the first import.
diff --git a/numpy/core/src/multiarray/arrayobject.c b/numpy/core/src/multiarray/arrayobject.c
@@ -62,6 +62,9 @@ maintainer email:  oliphant.travis@ieee.org
 #include "binop_override.h"
 #include "array_coercion.h"
 
+
+NPY_NO_EXPORT npy_bool numpy_warn_if_no_mem_policy = 0;
+
 /*NUMPY_API
   Compute the size of an array (in number of items)
 */
@@ -460,9 +463,8 @@ array_dealloc(PyArrayObject *self)
             }
         }
         if (fa->mem_handler == NULL) {
-            char *env = getenv("NUMPY_WARN_IF_NO_MEM_POLICY");
-            if ((env != NULL) && (strncmp(env, "1", 1) == 0)) {
-                char const * msg = "Trying to dealloc data, but a memory policy "
+            if (numpy_warn_if_no_mem_policy) {
+                char const *msg = "Trying to dealloc data, but a memory policy "
                     "is not set. If you take ownership of the data, you must "
                     "set a base owning the data (e.g. a PyCapsule).";
                 WARN_IN_DEALLOC(PyExc_RuntimeWarning, msg);

diff --git a/numpy/core/src/multiarray/arrayobject.h b/numpy/core/src/multiarray/arrayobject.h
@@ -5,6 +5,8 @@
 #ifndef NUMPY_CORE_SRC_MULTIARRAY_ARRAYOBJECT_H_
 #define NUMPY_CORE_SRC_MULTIARRAY_ARRAYOBJECT_H_
 
+extern NPY_NO_EXPORT npy_bool numpy_warn_if_no_mem_policy;
+
 NPY_NO_EXPORT PyObject *
 _strings_richcompare(PyArrayObject *self, PyArrayObject *other, int cmp_op,
                      int rstrip);

diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
@@ -4378,6 +4378,24 @@ normalize_axis_index(PyObject *NPY_UNUSED(self),
 }
 
 
+static PyObject *
+_set_numpy_warn_if_no_mem_policy(PyObject *NPY_UNUSED(self), PyObject *arg)
+{
+    int res = PyObject_IsTrue(arg);
+    if (res < 0) {
+        return NULL;
+    }
+    int old_value = numpy_warn_if_no_mem_policy;
+    numpy_warn_if_no_mem_policy = res;
+    if (old_value) {
+        Py_RETURN_TRUE;
+    }
+    else {
+        Py_RETURN_FALSE;
+    }
+}
+
+
 static PyObject *
 _reload_guard(PyObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args)) {
     static int initialized = 0;
@@ -4625,6 +4643,9 @@ static struct PyMethodDef array_module_methods[] = {
          METH_O, "Set the NEP 50 promotion state.  This is not thread-safe.\n"
                  "The optional warnings can be safely silenced using the \n"
                  "`np._no_nep50_warning()` context manager."},
+    {"_set_numpy_warn_if_no_mem_policy",
+         (PyCFunction)_set_numpy_warn_if_no_mem_policy,
+         METH_O, "Change the warn if no mem policy flag for testing."},
     {"_add_newdoc_ufunc", (PyCFunction)add_newdoc_ufunc,
         METH_VARARGS, NULL},
     {"_get_sfloat_dtype",
@@ -4913,6 +4934,14 @@ initialize_static_globals(void)
         return -1;
     }
 
+    char *env = getenv("NUMPY_WARN_IF_NO_MEM_POLICY");
+    if ((env != NULL) && (strncmp(env, "1", 1) == 0)) {
+        numpy_warn_if_no_mem_policy = 1;
+        }
+    else {
+        numpy_warn_if_no_mem_policy = 0;
+    }
+
     return 0;
 }
 

diff --git a/numpy/core/tests/test_mem_policy.py b/numpy/core/tests/test_mem_policy.py
@@ -409,16 +409,19 @@ def test_switch_owner(get_module, policy):
     a = get_module.get_array()
     assert np.core.multiarray.get_handler_name(a) is None
     get_module.set_own(a)
-    oldval = os.environ.get('NUMPY_WARN_IF_NO_MEM_POLICY', None)
+
     if policy is None:
-        if 'NUMPY_WARN_IF_NO_MEM_POLICY' in os.environ:
-            os.environ.pop('NUMPY_WARN_IF_NO_MEM_POLICY')
+        # See what we expect to be set based on the env variable
+        policy = os.getenv("NUMPY_WARN_IF_NO_MEM_POLICY", "0") == "1"
+        oldval = None
     else:
-        os.environ['NUMPY_WARN_IF_NO_MEM_POLICY'] = policy
+        policy = policy == "1"
+        oldval = np.core._multiarray_umath._set_numpy_warn_if_no_mem_policy(
+            policy)
     try:
         # The policy should be NULL, so we have to assume we can call
         # "free".  A warning is given if the policy == "1"
-        if policy == "1":
+        if policy:
             with assert_warns(RuntimeWarning) as w:
                 del a
                 gc.collect()
@@ -427,11 +430,8 @@ def test_switch_owner(get_module, policy):
             gc.collect()
 
     finally:
-        if oldval is None:
-            if 'NUMPY_WARN_IF_NO_MEM_POLICY' in os.environ:
-                os.environ.pop('NUMPY_WARN_IF_NO_MEM_POLICY')
-        else:
-            os.environ['NUMPY_WARN_IF_NO_MEM_POLICY'] = oldval
+        if oldval is not None:
+            np.core._multiarray_umath._set_numpy_warn_if_no_mem_policy(oldval)
 
 
 @pytest.mark.skipif(sys.version_info >= (3, 12), reason="no numpy.distutils")