Merge pull request #20238 from seberg/fixup-transfer-ownership

BUG: Fix environment checking logic for `NUMPY_WARN_IF_NO_MEM_POLICY`

2 files changed, 21 insertions, 6 deletions

diff --git a/numpy/core/src/multiarray/arrayobject.c b/numpy/core/src/multiarray/arrayobject.c
index 237fb9b72..c8aaced4e 100644
--- a/numpy/core/src/multiarray/arrayobject.c
+++ b/numpy/core/src/multiarray/arrayobject.c
@@ -503,7 +503,7 @@ 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)) {
+            if ((env != NULL) && (strncmp(env, "1", 1) == 0)) {
                 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).";
diff --git a/numpy/core/tests/test_mem_policy.py b/numpy/core/tests/test_mem_policy.py
index 4ec9b3b11..7fec8897f 100644
--- a/numpy/core/tests/test_mem_policy.py
+++ b/numpy/core/tests/test_mem_policy.py
@@ -357,24 +357,39 @@ def test_new_policy(get_module):
     c = np.arange(10)
     assert np.core.multiarray.get_handler_name(c) == orig_policy_name
 
-def test_switch_owner(get_module):
+@pytest.mark.xfail(sys.implementation.name == "pypy",
+                   reason=("bad interaction between getenv and "
+                           "os.environ inside pytest"))
+@pytest.mark.parametrize("policy", ["0", "1", None])
+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)
-    os.environ['NUMPY_WARN_IF_NO_MEM_POLICY'] = "1"
+    if policy 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'] = policy
     try:
         # The policy should be NULL, so we have to assume we can call
-        # "free"
-        with assert_warns(RuntimeWarning) as w:
+        # "free".  A warning is given if the policy == "1"
+        if policy == "1":
+            with assert_warns(RuntimeWarning) as w:
+                del a
+                gc.collect()
+        else:
             del a
             gc.collect()
+
     finally:
         if oldval is None:
-            os.environ.pop('NUMPY_WARN_IF_NO_MEM_POLICY')
+            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
 
+def test_owner_is_base(get_module):
     a = get_module.get_array_with_base()
     with pytest.warns(UserWarning, match='warn_on_free'):
         del a