ENH: add __array_function__ protocol in polynomial (#28996)

lvllvl · web-flow · commit 06f4d5ed1b85 · 2025-05-23T17:48:01.000+02:00
diff --git a/numpy/polynomial/polynomial.py b/numpy/polynomial/polynomial.py
@@ -82,6 +82,7 @@
 
 import numpy as np
 import numpy.linalg as la
+from numpy._core.overrides import array_function_dispatch
 from numpy.lib.array_utils import normalize_axis_index
 
 from . import polyutils as pu
@@ -841,7 +842,13 @@ def polyvalfromroots(x, r, tensor=True):
             raise ValueError("x.ndim must be < r.ndim when tensor == False")
     return np.prod(x - r, axis=0)
 
+def _polyval2d_dispatcher(x, y, c):
+    return (x, y, c)
 
+def _polygrid2d_dispatcher(x, y, c):
+    return (x, y, c)
+
+@array_function_dispatch(_polyval2d_dispatcher)
 def polyval2d(x, y, c):
     """
     Evaluate a 2-D polynomial at points (x, y).
@@ -893,7 +900,7 @@ def polyval2d(x, y, c):
     """
     return pu._valnd(polyval, c, x, y)
 
-
+@array_function_dispatch(_polygrid2d_dispatcher)
 def polygrid2d(x, y, c):
     """
     Evaluate a 2-D polynomial on the Cartesian product of x and y.
diff --git a/numpy/polynomial/tests/test_polynomial.py b/numpy/polynomial/tests/test_polynomial.py
@@ -667,3 +667,25 @@ def test_result_type(self):
 
         arr = np.polydiv(1, np.float32(1))
         assert_equal(arr[0].dtype, np.float64)
+
+class ArrayFunctionInterceptor:
+    def __init__(self):
+        self.called = False
+
+    def __array_function__(self, func, types, args, kwargs):
+        self.called = True
+        return "intercepted"
+
+def test_polyval2d_array_function_hook():
+    x = ArrayFunctionInterceptor()
+    y = ArrayFunctionInterceptor()
+    c = ArrayFunctionInterceptor()
+    result = np.polynomial.polynomial.polyval2d(x, y, c)
+    assert result == "intercepted"
+
+def test_polygrid2d_array_function_hook():
+    x = ArrayFunctionInterceptor()
+    y = ArrayFunctionInterceptor()
+    c = ArrayFunctionInterceptor()
+    result = np.polynomial.polynomial.polygrid2d(x, y, c)
+    assert result == "intercepted"
