TYP: Accept dispatcher function with optional returns in _core.overrides (numpy#29171)

jorenham · seberg · web-flow · commit 29c85f376bc4 · 2025-06-11T19:58:40.000+02:00
Co-authored-by: Sebastian Berg &lt;sebastianb@nvidia.com&gt;
diff --git a/numpy/_core/overrides.pyi b/numpy/_core/overrides.pyi
@@ -1,11 +1,11 @@
 from collections.abc import Callable, Iterable
-from typing import Any, Final, NamedTuple, ParamSpec, TypeVar
-
-from numpy._typing import _SupportsArrayFunc
+from typing import Any, Final, NamedTuple, ParamSpec, TypeAlias, TypeVar
 
 _T = TypeVar("_T")
 _Tss = ParamSpec("_Tss")
-_FuncT = TypeVar("_FuncT", bound=Callable[..., object])
+_FuncLikeT = TypeVar("_FuncLikeT", bound=type | Callable[..., object])
+
+_Dispatcher: TypeAlias = Callable[_Tss, Iterable[object]]
 
 ###
 
@@ -18,31 +18,28 @@ class ArgSpec(NamedTuple):
     keywords: str | None
     defaults: tuple[Any, ...]
 
-def get_array_function_like_doc(public_api: Callable[..., Any], docstring_template: str = "") -> str: ...
-def finalize_array_function_like(public_api: _FuncT) -> _FuncT: ...
+def get_array_function_like_doc(public_api: Callable[..., object], docstring_template: str = "") -> str: ...
+def finalize_array_function_like(public_api: _FuncLikeT) -> _FuncLikeT: ...
 
 #
-def verify_matching_signatures(
-    implementation: Callable[_Tss, object],
-    dispatcher: Callable[_Tss, Iterable[_SupportsArrayFunc]],
-) -> None: ...
+def verify_matching_signatures(implementation: Callable[_Tss, object], dispatcher: _Dispatcher[_Tss]) -> None: ...
 
 # NOTE: This actually returns a `_ArrayFunctionDispatcher` callable wrapper object, with
 # the original wrapped callable stored in the `._implementation` attribute. It checks
 # for any `__array_function__` of the values of specific arguments that the dispatcher
 # specifies. Since the dispatcher only returns an iterable of passed array-like args,
 # this overridable behaviour is impossible to annotate.
 def array_function_dispatch(
-    dispatcher: Callable[_Tss, Iterable[_SupportsArrayFunc]] | None = None,
+    dispatcher: _Dispatcher[_Tss] | None = None,
     module: str | None = None,
     verify: bool = True,
     docs_from_dispatcher: bool = False,
-) -> Callable[[_FuncT], _FuncT]: ...
+) -> Callable[[_FuncLikeT], _FuncLikeT]: ...
 
 #
 def array_function_from_dispatcher(
     implementation: Callable[_Tss, _T],
     module: str | None = None,
     verify: bool = True,
     docs_from_dispatcher: bool = True,
-) -> Callable[[Callable[_Tss, Iterable[_SupportsArrayFunc]]], Callable[_Tss, _T]]: ...
+) -> Callable[[_Dispatcher[_Tss]], Callable[_Tss, _T]]: ...
