ENH: add __array_function__ protocol in polynomial

lvllvl · lvllvl · commit b5737e837500 · 2025-05-17T20:28:52.000Z
diff --git a/numpy/polynomial/polynomial.py b/numpy/polynomial/polynomial.py
@@ -83,6 +83,7 @@
 import numpy as np
 import numpy.linalg as la
 from numpy.lib.array_utils import normalize_axis_index
+from numpy._core.overrides import array_function_dispatch
 
 from . import polyutils as pu
 from ._polybase import ABCPolyBase
@@ -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"
