TYP: Add type tests for the numpy.polynomial series-specific methods

jorenham · jorenham · commit c7998c2930a4 · 2024-07-17T01:24:09.000+02:00
diff --git a/numpy/polynomial/chebyshev.pyi b/numpy/polynomial/chebyshev.pyi
@@ -86,11 +86,11 @@ def _zseries_to_cseries(zs: npt.NDArray[_SCT]) -> _Series[_SCT]: ...
 def _zseries_mul(
     z1: npt.NDArray[_SCT],
     z2: npt.NDArray[_SCT],
-) -> _Series[_SCT]:    ...
+) -> _Series[_SCT]: ...
 def _zseries_div(
     z1: npt.NDArray[_SCT],
     z2: npt.NDArray[_SCT],
-) -> _Series[_SCT]:    ...
+) -> _Series[_SCT]: ...
 def _zseries_der(zs: npt.NDArray[_SCT]) -> _Series[_SCT]: ...
 def _zseries_int(zs: npt.NDArray[_SCT]) -> _Series[_SCT]: ...
 
@@ -130,27 +130,25 @@ chebpts1: _FuncPts[L["chebpts1"]]
 chebpts2: _FuncPts[L["chebpts2"]]
 
 # keep in sync with `Chebyshev.interpolate`
-_RT = TypeVar(
-    "_RT",
-    np.float64,
-    np.complex128,
-    np.floating[Any],
-    np.complexfloating[Any, Any],
-    np.number[Any],
-    np.object_,
-)
+_RT = TypeVar("_RT", bound=np.number[Any] | np.bool | np.object_)
+@overload
+def chebinterpolate(
+    func: np.ufunc,
+    deg: _IntLike_co,
+    args: tuple[()] = ...,
+) -> npt.NDArray[np.float64 | np.complex128 | np.object_]: ...
 @overload
 def chebinterpolate(
     func: Callable[[npt.NDArray[np.float64]], _RT],
     deg: _IntLike_co,
     args: tuple[()] = ...,
-) -> npt.NDArray[_RT]:    ...
+) -> npt.NDArray[_RT]: ...
 @overload
 def chebinterpolate(
     func: Callable[Concatenate[npt.NDArray[np.float64], ...], _RT],
     deg: _IntLike_co,
     args: Iterable[Any],
-) -> npt.NDArray[_RT]:    ...
+) -> npt.NDArray[_RT]: ...
 
 _Self = TypeVar("_Self", bound=object)
 
@@ -164,7 +162,7 @@ class Chebyshev(ABCPolyBase[L["T"]]):
         deg: _IntLike_co,
         domain: None | _SeriesLikeCoef_co = ...,
         args: tuple[()] = ...,
-    ) -> _Self:        ...
+    ) -> _Self: ...
     @overload
     @classmethod
     def interpolate(
@@ -178,7 +176,7 @@ class Chebyshev(ABCPolyBase[L["T"]]):
         domain: None | _SeriesLikeCoef_co = ...,
         *,
         args: Iterable[Any],
-    ) -> _Self:        ...
+    ) -> _Self: ...
     @overload
     @classmethod
     def interpolate(
@@ -191,4 +189,4 @@ class Chebyshev(ABCPolyBase[L["T"]]):
         domain: None | _SeriesLikeCoef_co,
         args: Iterable[Any],
         /,
-    ) -> _Self:        ...
+    ) -> _Self: ...
diff --git a/numpy/typing/tests/data/reveal/polynomial_series.pyi b/numpy/typing/tests/data/reveal/polynomial_series.pyi
@@ -25,6 +25,7 @@ AR_c16: npt.NDArray[np.complex128]
 AR_O: npt.NDArray[np.object_]
 
 PS_poly: npp.Polynomial
+PS_cheb: npp.Chebyshev
 
 assert_type(npp.polynomial.polyroots(AR_f8), _ArrFloat1D64)
 assert_type(npp.polynomial.polyroots(AR_c16), _ArrComplex1D128)
@@ -119,3 +120,25 @@ assert_type(
     npp.polynomial.polyfit(AR_f8, AR_c16, 1, full=True)[0],
     npt.NDArray[np.complexfloating[Any, Any]],
 )
+
+assert_type(npp.chebyshev.chebgauss(2), tuple[_ArrFloat1D64, _ArrFloat1D64])
+
+assert_type(npp.chebyshev.chebweight(AR_f8), npt.NDArray[np.float64])
+assert_type(npp.chebyshev.chebweight(AR_c16), npt.NDArray[np.complex128])
+assert_type(npp.chebyshev.chebweight(AR_O), npt.NDArray[np.object_])
+
+assert_type(npp.chebyshev.poly2cheb(AR_f8), _ArrFloat1D)
+assert_type(npp.chebyshev.poly2cheb(AR_c16), _ArrComplex1D)
+assert_type(npp.chebyshev.poly2cheb(AR_O), _ArrObject1D)
+
+assert_type(npp.chebyshev.cheb2poly(AR_f8), _ArrFloat1D)
+assert_type(npp.chebyshev.cheb2poly(AR_c16), _ArrComplex1D)
+assert_type(npp.chebyshev.cheb2poly(AR_O), _ArrObject1D)
+
+assert_type(npp.chebyshev.chebpts1(6), _ArrFloat1D64)
+assert_type(npp.chebyshev.chebpts2(6), _ArrFloat1D64)
+
+assert_type(
+    npp.chebyshev.chebinterpolate(np.tanh, 3),
+    npt.NDArray[np.float64 | np.complex128 | np.object_],
+)
