BUG: Fix incorrect GIL release in array.nonzero (#13930)

Added checks to array.nonzero to ensure GIL is not released if
it is needed.

Closes gh-13753

2 files changed, 58 insertions, 2 deletions

diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c
index 9351b5fc7..762563eb5 100644
--- a/numpy/core/src/multiarray/item_selection.c
+++ b/numpy/core/src/multiarray/item_selection.c
@@ -2203,9 +2203,13 @@ PyArray_Nonzero(PyArrayObject *self)
     PyArrayObject *ret = NULL;
     PyObject *ret_tuple;
     npy_intp ret_dims[2];
-    PyArray_NonzeroFunc *nonzero = PyArray_DESCR(self)->f->nonzero;
+
+    PyArray_NonzeroFunc *nonzero;
+    PyArray_Descr *dtype;
+
     npy_intp nonzero_count;
     npy_intp added_count = 0;
+    int needs_api;
     int is_bool;
 
     NpyIter *iter;
@@ -2213,6 +2217,10 @@ PyArray_Nonzero(PyArrayObject *self)
     NpyIter_GetMultiIndexFunc *get_multi_index;
     char **dataptr;
 
+    dtype = PyArray_DESCR(self);
+    nonzero = dtype->f->nonzero;
+    needs_api = PyDataType_FLAGCHK(dtype, NPY_NEEDS_PYAPI);
+
     /* Special case - nonzero(zero_d) is nonzero(atleast_1d(zero_d)) */
     if (ndim == 0) {
         char const* msg;
@@ -2283,7 +2291,9 @@ PyArray_Nonzero(PyArrayObject *self)
             goto finish;
         }
 
-        NPY_BEGIN_THREADS_THRESHOLDED(count);
+        if (!needs_api) {
+            NPY_BEGIN_THREADS_THRESHOLDED(count);
+        }
 
         /* avoid function call for bool */
         if (is_bool) {
@@ -2324,6 +2334,9 @@ PyArray_Nonzero(PyArrayObject *self)
                     }
                     *multi_index++ = j;
                 }
+                if (needs_api && PyErr_Occurred()) {
+                    break;
+                }
                 data += stride;
             }
         }
@@ -2364,6 +2377,8 @@ PyArray_Nonzero(PyArrayObject *self)
             Py_DECREF(ret);
             return NULL;
         }
+        
+        needs_api = NpyIter_IterationNeedsAPI(iter);
 
         NPY_BEGIN_THREADS_NDITER(iter);
 
@@ -2390,6 +2405,9 @@ PyArray_Nonzero(PyArrayObject *self)
                     get_multi_index(iter, multi_index);
                     multi_index += ndim;
                 }
+                if (needs_api && PyErr_Occurred()) {
+                    break;
+                }
             } while(iternext(iter));
         }
 
@@ -2399,6 +2417,11 @@ PyArray_Nonzero(PyArrayObject *self)
     NpyIter_Deallocate(iter);
 
 finish:
+    if (PyErr_Occurred()) {
+        Py_DECREF(ret);
+        return NULL;
+    }
+    
     /* if executed `nonzero()` check for miscount due to side-effect */
     if (!is_bool && added_count != nonzero_count) {
         PyErr_SetString(PyExc_RuntimeError,
diff --git a/numpy/core/tests/test_numeric.py b/numpy/core/tests/test_numeric.py
index 935b84234..3e85054b7 100644
--- a/numpy/core/tests/test_numeric.py
+++ b/numpy/core/tests/test_numeric.py
@@ -1254,6 +1254,39 @@ class TestNonzero(object):
         a = np.array([[False], [TrueThenFalse()]])
         assert_raises(RuntimeError, np.nonzero, a)
 
+    def test_nonzero_exception_safe(self):
+        # gh-13930
+
+        class ThrowsAfter:
+            def __init__(self, iters):
+                self.iters_left = iters
+
+            def __bool__(self):
+                if self.iters_left == 0:
+                    raise ValueError("called `iters` times")
+
+                self.iters_left -= 1
+                return True
+
+        """
+        Test that a ValueError is raised instead of a SystemError
+
+        If the __bool__ function is called after the error state is set,
+        Python (cpython) will raise a SystemError.
+        """
+
+        # assert that an exception in first pass is handled correctly
+        a = np.array([ThrowsAfter(5)]*10)
+        assert_raises(ValueError, np.nonzero, a)
+
+        # raise exception in second pass for 1-dimensional loop
+        a = np.array([ThrowsAfter(15)]*10)
+        assert_raises(ValueError, np.nonzero, a)
+
+        # raise exception in second pass for n-dimensional loop
+        a = np.array([[ThrowsAfter(15)]]*10)
+        assert_raises(ValueError, np.nonzero, a)
+
 
 class TestIndex(object):
     def test_boolean(self):