lint

Author: crusaderky

src/array_api_extra/_apply.py

@@ -17,43 +17,43 @@
 
 if TYPE_CHECKING:
     # https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
-    from typing import TypeAlias
+    from typing import ParamSpec, TypeAlias
 
     import numpy as np
 
     NumPyObject: TypeAlias = np.ndarray[Any, Any] | np.generic  # type: ignore[no-any-explicit]
-    KwArg: TypeAlias = Any  # type: ignore[no-any-explicit]
+    P = ParamSpec("P")
 
 
 @overload
-def apply_numpy_func(
-    func: Callable[..., NumPyObject],
+def apply_numpy_func(  # type: ignore[valid-type]
+    func: Callable[P, NumPyObject],
     *args: Array,
     shape: tuple[int, ...] | None = None,
     dtype: DType | None = None,
     xp: ModuleType | None = None,
-    **kwargs: KwArg,
+    **kwargs: P.kwargs,  # pyright: ignore[reportGeneralTypeIssues]
 ) -> Array: ...  # numpydoc ignore=GL08
 
 
 @overload
-def apply_numpy_func(  # type: ignore[no-any-decorated]
-    func: Callable[..., Sequence[NumPyObject]],
+def apply_numpy_func(  # type: ignore[valid-type]
+    func: Callable[P, Sequence[NumPyObject]],
     *args: Array,
     shape: Sequence[tuple[int, ...]],
     dtype: Sequence[DType] | None = None,
     xp: ModuleType | None = None,
-    **kwargs: Any,
+    **kwargs: P.kwargs,  # pyright: ignore[reportGeneralTypeIssues]
 ) -> tuple[Array, ...]: ...  # numpydoc ignore=GL08
 
 
-def apply_numpy_func(  # type: ignore[no-any-explicit]
-    func: Callable[..., NumPyObject | Sequence[NumPyObject]],
+def apply_numpy_func(  # type: ignore[valid-type]
+    func: Callable[P, NumPyObject | Sequence[NumPyObject]],
     *args: Array,
     shape: tuple[int, ...] | Sequence[tuple[int, ...]] | None = None,
     dtype: DType | Sequence[DType] | None = None,
     xp: ModuleType | None = None,
-    **kwargs: Any,
+    **kwargs: P.kwargs,  # pyright: ignore[reportGeneralTypeIssues]
 ) -> Array | tuple[Array, ...]:
     """
     Apply a function that operates on NumPy arrays to Array API compliant arrays.
@@ -139,7 +139,7 @@ def apply_numpy_func(  # type: ignore[no-any-explicit]
     elif isinstance(shape, tuple) and all(isinstance(s, int) for s in shape):
         shapes = [shape]
     else:
-        shapes = shape
+        shapes = list(shape)
         multi_output = True
 
     if dtype is None:
@@ -148,7 +148,7 @@ def apply_numpy_func(  # type: ignore[no-any-explicit]
         if not isinstance(dtype, Sequence):
             msg = "Got sequence of shapes but only one dtype"
             raise TypeError(msg)
-        dtypes = dtype
+        dtypes = list(dtype)  # pyright: ignore[reportUnknownArgumentType]
     else:
         if isinstance(dtype, Sequence):
             msg = "Got single shape but multiple dtypes"
@@ -254,13 +254,16 @@ def wrapper(  # type: ignore[no-any-decorated,no-any-explicit]
         args = tuple(np.asarray(arg) for arg in args)
         out = func(*args, **kwargs)
 
+        # Stay relaxed on output validation, e.g. in case func returns a
+        # Python scalar instead of a np.generic
         if multi_output:
             if not isinstance(out, Sequence) or isinstance(out, np.ndarray):
                 msg = "Expected multiple outputs, got a single one"
                 raise ValueError(msg)
+            outs = out
         else:
-            out = (out,)
+            outs = [cast("NumPyObject", out)]
 
-        return tuple(xp.asarray(o) for o in out)
+        return tuple(xp.asarray(o) for o in outs)
 
     return wrapper