TYP: Improved numpy.piecewise type-hints with ParamSpec

Author: jorenham

numpy/lib/_function_base_impl.pyi

@@ -1,7 +1,9 @@
 from collections.abc import Sequence, Iterator, Callable, Iterable
 from typing import (
+    Concatenate,
     Literal as L,
     Any,
+    ParamSpec,
     TypeVar,
     overload,
     Protocol,
@@ -34,6 +36,7 @@ from numpy._typing import (
     _ScalarLike_co,
     _DTypeLike,
     _ArrayLike,
+    _ArrayLikeBool_co,
     _ArrayLikeInt_co,
     _ArrayLikeFloat_co,
     _ArrayLikeComplex_co,
@@ -50,6 +53,8 @@ from numpy._core.multiarray import (
 
 _T = TypeVar("_T")
 _T_co = TypeVar("_T_co", covariant=True)
+# The `{}ss` suffix refers to the Python 3.12 syntax: `**P`
+_Pss = ParamSpec("_Pss")
 _SCT = TypeVar("_SCT", bound=generic)
 _ArrayType = TypeVar("_ArrayType", bound=NDArray[Any])
 
@@ -180,23 +185,29 @@ def asarray_chkfinite(
     order: _OrderKACF = ...,
 ) -> NDArray[Any]: ...
 
-# TODO: Use PEP 612 `ParamSpec` once mypy supports `Concatenate`
-# xref python/mypy#8645
 @overload
 def piecewise(
     x: _ArrayLike[_SCT],
-    condlist: ArrayLike,
-    funclist: Sequence[Any | Callable[..., Any]],
-    *args: Any,
-    **kw: Any,
+    condlist: _ArrayLike[bool_] | Sequence[_ArrayLikeBool_co],
+    funclist: Sequence[
+        Callable[Concatenate[NDArray[_SCT], _Pss], NDArray[_SCT | Any]]
+        | _SCT | object
+    ],
+    /,
+    *args: _Pss.args,
+    **kw: _Pss.kwargs,
 ) -> NDArray[_SCT]: ...
 @overload
 def piecewise(
     x: ArrayLike,
-    condlist: ArrayLike,
-    funclist: Sequence[Any | Callable[..., Any]],
-    *args: Any,
-    **kw: Any,
+    condlist: _ArrayLike[bool_] | Sequence[_ArrayLikeBool_co],
+    funclist: Sequence[
+        Callable[Concatenate[NDArray[Any], _Pss], NDArray[Any]]
+        | object
+    ],
+    /,
+    *args: _Pss.args,
+    **kw: _Pss.kwargs,
 ) -> NDArray[Any]: ...
 
 def select(

numpy/typing/tests/data/fail/lib_function_base.pyi

@@ -8,8 +8,10 @@ AR_c16: npt.NDArray[np.complex128]
 AR_m: npt.NDArray[np.timedelta64]
 AR_M: npt.NDArray[np.datetime64]
 AR_O: npt.NDArray[np.object_]
+AR_b_list: list[npt.NDArray[np.bool]]
 
-def func(a: int) -> None: ...
+def fn_none_i(a: None, /) -> npt.NDArray[Any]: ...
+def fn_ar_i(a: npt.NDArray[np.float64], posarg: int, /) -> npt.NDArray[Any]: ...
 
 np.average(AR_m)  # E: incompatible type
 np.select(1, [AR_f8])  # E: incompatible type
@@ -21,6 +23,15 @@ np.place(1, [True], 1.5)  # E: incompatible type
 np.vectorize(1)  # E: incompatible type
 np.place(AR_f8, slice(None), 5)  # E: incompatible type
 
+np.piecewise(AR_f8, True, [fn_ar_i], 42)  # E: No overload variants
+# TODO: enable these once mypy actually supports ParamSpec (released in 2021)
+# NOTE: pyright correctly reports errors for these (`reportCallIssue`)
+# np.piecewise(AR_f8, AR_b_list, [fn_none_i])  # E: No overload variants
+# np.piecewise(AR_f8, AR_b_list, [fn_ar_i])  # E: No overload variant
+# np.piecewise(AR_f8, AR_b_list, [fn_ar_i], 3.14)  # E: No overload variant
+# np.piecewise(AR_f8, AR_b_list, [fn_ar_i], 42, None)  # E: No overload variant
+# np.piecewise(AR_f8, AR_b_list, [fn_ar_i], 42, _=None)  # E: No overload variant
+
 np.interp(AR_f8, AR_c16, AR_f8)  # E: incompatible type
 np.interp(AR_c16, AR_f8, AR_f8)  # E: incompatible type
 np.interp(AR_f8, AR_f8, AR_f8, period=AR_c16)  # E: No overload variant

numpy/typing/tests/data/reveal/lib_function_base.pyi

@@ -28,7 +28,16 @@ AR_b: npt.NDArray[np.bool]
 AR_U: npt.NDArray[np.str_]
 CHAR_AR_U: np.char.chararray[Any, np.dtype[np.str_]]
 
-def func(*args: Any, **kwargs: Any) -> Any: ...
+AR_b_list: list[npt.NDArray[np.bool]]
+
+def func(
+    a: npt.NDArray[Any],
+    posarg: bool = ...,
+    /,
+    arg: int = ...,
+    *,
+    kwarg: str = ...,
+) -> npt.NDArray[Any]: ...
 
 assert_type(vectorized_func.pyfunc, Callable[..., Any])
 assert_type(vectorized_func.cache, bool)
@@ -69,7 +78,10 @@ assert_type(np.asarray_chkfinite(AR_f8, dtype=np.float64), npt.NDArray[np.float6
 assert_type(np.asarray_chkfinite(AR_f8, dtype=float), npt.NDArray[Any])
 
 assert_type(np.piecewise(AR_f8, AR_b, [func]), npt.NDArray[np.float64])
-assert_type(np.piecewise(AR_LIKE_f8, AR_b, [func]), npt.NDArray[Any])
+assert_type(np.piecewise(AR_f8, AR_b_list, [func]), npt.NDArray[np.float64])
+assert_type(np.piecewise(AR_f8, AR_b_list, [func], True, -1, kwarg=''), npt.NDArray[np.float64])
+assert_type(np.piecewise(AR_f8, AR_b_list, [func], True, arg=-1, kwarg=''), npt.NDArray[np.float64])
+assert_type(np.piecewise(AR_LIKE_f8, AR_b_list, [func]), npt.NDArray[Any])
 
 assert_type(np.select([AR_f8], [AR_f8]), npt.NDArray[Any])