TST: Add tests for PyUFunc_ReplaceLoopBySignature (numpy#31097)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Sebastian Berg <sebastian@sipsolutions.net>

Author: 3 people

numpy/_core/_umath_tests.pyi

@@ -30,6 +30,10 @@ def test_signature(
 # undocumented
 def test_dispatch() -> _TestDispatchResult: ...
 
+# undocumented test helpers for PyUFunc_ReplaceLoopBySignature
+def replace_loop(ufunc: np.ufunc, /) -> object: ...
+def restore_loop(ufunc: np.ufunc, capsule: object, /) -> None: ...
+
 # undocumented ufuncs and gufuncs
 always_error: Final[_UFunc_Nin2_Nout1[L["always_error"], L[1], None]] = ...
 always_error_unary: Final[_UFunc_Nin1_Nout1[L["always_error_unary"], L[1], None]] = ...

numpy/_core/src/umath/_umath_tests.c.src

@@ -857,6 +857,72 @@ static void *const conv1d_full_data[] = {NULL};
 static const char conv1d_full_typecodes[] = {NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE};
 
 
+/*
+ * Test helpers for PyUFunc_ReplaceLoopBySignature.
+ *
+ * _constant42_loop: unary float64 loop that always writes 42.0.
+ */
+static void
+_constant42_loop(char **args, npy_intp const *dimensions,
+                 npy_intp const *steps, void *NPY_UNUSED(data))
+{
+    npy_intp n = dimensions[0];
+    char *out = args[1];
+    npy_intp out_step = steps[1];
+    for (npy_intp i = 0; i < n; i++) {
+        *(double *)out = 42.0;
+        out += out_step;
+    }
+}
+
+/*
+ * replace_loop(ufunc): Replace the dd->d loop with _constant42_loop.
+ * Only works for unary ufuncs.  Returns a capsule holding the old loop.
+ */
+static PyObject *
+UMath_Tests_replace_loop(PyObject *NPY_UNUSED(dummy), PyObject *args)
+{
+    PyUFuncObject *ufunc;
+    if (!PyArg_ParseTuple(args, "O!", &PyUFunc_Type, &ufunc)) {
+        return NULL;
+    }
+    if (ufunc->nin != 1 || ufunc->nout != 1) {
+        PyErr_SetString(PyExc_ValueError,
+                "replace_loop only supports unary ufuncs");
+        return NULL;
+    }
+    int signature[2] = {NPY_DOUBLE, NPY_DOUBLE};
+    PyUFuncGenericFunction oldfunc = NULL;
+    if (PyUFunc_ReplaceLoopBySignature(
+            ufunc, _constant42_loop, signature, &oldfunc) < 0) {
+        PyErr_SetString(PyExc_RuntimeError,
+                "failed to find a float64 loop");
+        return NULL;
+    }
+    return PyCapsule_New((void *)oldfunc, "oldfunc", NULL);
+}
+
+/*
+ * restore_loop(ufunc, capsule): Restore the loop saved by replace_loop.
+ */
+static PyObject *
+UMath_Tests_restore_loop(PyObject *NPY_UNUSED(dummy), PyObject *args)
+{
+    PyUFuncObject *ufunc;
+    PyObject *capsule;
+    if (!PyArg_ParseTuple(args, "O!O", &PyUFunc_Type, &ufunc, &capsule)) {
+        return NULL;
+    }
+    PyUFuncGenericFunction oldfunc = (PyUFuncGenericFunction)
+            PyCapsule_GetPointer(capsule, "oldfunc");
+    if (oldfunc == NULL) {
+        return NULL;
+    }
+    int signature[2] = {NPY_DOUBLE, NPY_DOUBLE};
+    PyUFunc_ReplaceLoopBySignature(ufunc, oldfunc, signature, NULL);
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef UMath_TestsMethods[] = {
     {"test_signature",  UMath_Tests_test_signature, METH_VARARGS,
      "Test signature parsing of ufunc. \n"
@@ -865,6 +931,12 @@ static PyMethodDef UMath_TestsMethods[] = {
      "internals. \n",
      },
     {"test_dispatch", UMath_Tests_test_dispatch, METH_NOARGS, NULL},
+    {"replace_loop", UMath_Tests_replace_loop, METH_VARARGS,
+     "Replace the float64 loop of a ufunc with one that outputs 42.0.\n"
+     "Returns a capsule holding the old loop for restore_loop().\n"},
+    {"restore_loop", UMath_Tests_restore_loop, METH_VARARGS,
+     "Restore a ufunc loop previously replaced by replace_loop().\n"
+     "Arguments: ufunc, capsule\n"},
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };
 

numpy/_core/tests/test_umath.py

@@ -5111,6 +5111,26 @@ def test_bad_legacy_gufunc_silent_errors(x1):
         ncu_tests.always_error_gufunc(x1, 0.0)
 
 
+class TestReplaceLoopBySignature:
+    """Tests for PyUFunc_ReplaceLoopBySignature C API."""
+
+    @pytest.mark.thread_unsafe(reason="modifies ufunc within test")
+    def test_replace_loop(self):
+        # Call the ufunc first to populate any internal dispatch caches,
+        # then replace the float64 loop with one that outputs 42.0,
+        # verify the replacement is used, and restore the original.
+        a = np.array([1.0, 2.0, 3.0])
+        assert_array_equal(np.negative(a), [-1.0, -2.0, -3.0])
+
+        saved = ncu_tests.replace_loop(np.negative)
+        try:
+            assert_array_equal(np.negative(a), [42.0, 42.0, 42.0])
+        finally:
+            ncu_tests.restore_loop(np.negative, saved)
+
+        assert_array_equal(np.negative(a), [-1.0, -2.0, -3.0])
+
+
 class TestAddDocstring:
     @pytest.mark.skipif(sys.flags.optimize == 2, reason="Python running -OO")
     def test_add_same_docstring(self):