Merge pull request #5828 from jaimefrio/frompyfunc_segfault

BUG: Fix segfault in frompyfunc with nout=0

2 files changed, 23 insertions, 1 deletions

diff --git a/numpy/core/src/umath/loops.c.src b/numpy/core/src/umath/loops.c.src
index a69fc8147..e10482b06 100644
--- a/numpy/core/src/umath/loops.c.src
+++ b/numpy/core/src/umath/loops.c.src
@@ -616,11 +616,22 @@ PyUFunc_On_Om(char **args, npy_intp *dimensions, npy_intp *steps, void *func)
             }
             Py_DECREF(result);
         }
-        else {
+        else if (nout == 1) {
             op = (PyObject **)ptrs[nin];
             Py_XDECREF(*op);
             *op = result;
         }
+        else {
+            /*
+             * The single output does not match the expected nout != 1,
+             * which may be reasonable for nout == 0. For other values,
+             * behave as if a 1-tuple had been returned and exit.
+             */
+            Py_DECREF(result);
+            if (nout != 0) {
+                return;
+            }
+        }
         for(j = 0; j < ntot; j++) {
             ptrs[j] += steps[j];
         }
diff --git a/numpy/core/tests/test_regression.py b/numpy/core/tests/test_regression.py
index fa2f52a23..9e8511a01 100644
--- a/numpy/core/tests/test_regression.py
+++ b/numpy/core/tests/test_regression.py
@@ -2125,6 +2125,17 @@ class TestRegression(TestCase):
         for axis in chain(range(-a.ndim, a.ndim), [None]):
             assert_equal(a.repeat(2, axis=axis), a.repeat([2], axis=axis))
 
+    def test_frompyfunc_nout_0(self):
+        # gh-2014
+
+        def f(x):
+            x[0], x[-1] = x[-1], x[0]
+
+        uf = np.frompyfunc(f, 1, 0)
+        a = np.array([[1, 2, 3], [4, 5], [6, 7, 8, 9]])
+        assert_equal(uf(a), ())
+        assert_array_equal(a, [[3, 2, 1], [5, 4], [9, 7, 8, 6]])
+
 
 if __name__ == "__main__":
     run_module_suite()