Merge pull request #24255 from charris/backport-24248

BUG: Factor out slow `getenv` call used for memory policy warning

doc/source/reference/global_state.rst

@@ -96,3 +96,6 @@ release (NumPy 2.0) by setting an environment before importing NumPy:
 
 By default this will also activate the :ref:`NEP 50 <NEP50>` related setting
 ``NPY_PROMOTION_STATE`` (please see the NEP for details on this).
+
+.. versionchanged:: 1.25.2
+    This variable is only checked on the first import.

numpy/core/src/multiarray/arrayobject.c

@@ -62,6 +62,9 @@ maintainer email:  [email protected]
 #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);

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);

numpy/core/src/multiarray/multiarraymodule.c

@@ -4475,6 +4475,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;
@@ -4733,6 +4751,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",
@@ -5029,6 +5050,14 @@ initialize_static_globals(void)
         npy_numpy2_behavior = NPY_TRUE;
     }
 
+    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;
 }
 

numpy/core/tests/test_mem_policy.py

@@ -395,16 +395,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()
@@ -413,11 +416,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)
 
 def test_owner_is_base(get_module):
     a = get_module.get_array_with_base()