Merge pull request #24031 from seberg/array-ufunc-no-kwargs

BUG: Ensure __array_ufunc__ works without any kwargs passed
numpy · Jun 24, 2023 · 0e2af84 · 0e2af84
2 parents dd231b0 + 0538f9e
commit 0e2af84
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c
@@ -1135,6 +1135,9 @@ any_array_ufunc_overrides(PyObject *args, PyObject *kwds)
         }
     }
     Py_DECREF(fast);
+    if (kwds == NULL) {
+        return 0;
+    }
     /* check outputs, if any */
     nout = PyUFuncOverride_GetOutObjects(kwds, &out_kwd_obj, &out_objs);
     if (nout < 0) {

diff --git a/numpy/core/tests/test_umath.py b/numpy/core/tests/test_umath.py
@@ -3916,6 +3916,19 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
         assert_equal(a, check)
         assert_(a.info, {'inputs': [0, 2]})
 
+    def test_array_ufunc_direct_call(self):
+        # This is mainly a regression test for gh-24023 (shouldn't segfault)
+        a = np.array(1)
+        with pytest.raises(TypeError):
+            a.__array_ufunc__()
+
+        # No kwargs means kwargs may be NULL on the C-level
+        with pytest.raises(TypeError):
+            a.__array_ufunc__(1, 2)
+
+        # And the same with a valid call:
+        res = a.__array_ufunc__(np.add, "__call__", a, a)
+        assert_array_equal(res, a + a)
 
 class TestChoose:
     def test_mixed(self):