ENH: Add dtype-typing support to `np.core.function_base`

1 files changed, 162 insertions, 22 deletions

diff --git a/numpy/core/function_base.pyi b/numpy/core/function_base.pyi
index 68d3b3a98..6e0843a0e 100644
--- a/numpy/core/function_base.pyi
+++ b/numpy/core/function_base.pyi
@@ -1,55 +1,195 @@
-from typing import overload, Tuple, Union, Sequence, Any, SupportsIndex, Literal, List
+from typing import (
+    Literal as L,
+    overload,
+    Tuple,
+    Union,
+    Any,
+    SupportsIndex,
+    List,
+    Type,
+    TypeVar,
+)
 
-from numpy import ndarray
-from numpy.typing import ArrayLike, DTypeLike, _SupportsArray, _NumberLike_co
+from numpy import floating, complexfloating, generic, dtype
+from numpy.typing import (
+    NDArray,
+    ArrayLike,
+    DTypeLike,
+    _SupportsDType,
+    _SupportsArray,
+    _NumberLike_co,
+    _ArrayLikeFloat_co,
+    _ArrayLikeComplex_co,
+)
 
-# TODO: wait for support for recursive types
-_ArrayLikeNested = Sequence[Sequence[Any]]
-_ArrayLikeNumber = Union[
-    _NumberLike_co, Sequence[_NumberLike_co], ndarray, _SupportsArray, _ArrayLikeNested
+_SCT = TypeVar("_SCT", bound=generic)
+
+_DTypeLike = Union[
+    dtype[_SCT],
+    Type[_SCT],
+    _SupportsDType[dtype[_SCT]],
 ]
 
 __all__: List[str]
 
 @overload
 def linspace(
-    start: _ArrayLikeNumber,
-    stop: _ArrayLikeNumber,
+    start: _ArrayLikeFloat_co,
+    stop: _ArrayLikeFloat_co,
+    num: SupportsIndex = ...,
+    endpoint: bool = ...,
+    retstep: L[False] = ...,
+    dtype: None = ...,
+    axis: SupportsIndex = ...,
+) -> NDArray[floating[Any]]: ...
+@overload
+def linspace(
+    start: _ArrayLikeComplex_co,
+    stop: _ArrayLikeComplex_co,
+    num: SupportsIndex = ...,
+    endpoint: bool = ...,
+    retstep: L[False] = ...,
+    dtype: None = ...,
+    axis: SupportsIndex = ...,
+) -> NDArray[complexfloating[Any, Any]]: ...
+@overload
+def linspace(
+    start: _ArrayLikeComplex_co,
+    stop: _ArrayLikeComplex_co,
+    num: SupportsIndex = ...,
+    endpoint: bool = ...,
+    retstep: L[False] = ...,
+    dtype: _DTypeLike[_SCT] = ...,
+    axis: SupportsIndex = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def linspace(
+    start: _ArrayLikeComplex_co,
+    stop: _ArrayLikeComplex_co,
     num: SupportsIndex = ...,
     endpoint: bool = ...,
-    retstep: Literal[False] = ...,
+    retstep: L[False] = ...,
     dtype: DTypeLike = ...,
     axis: SupportsIndex = ...,
-) -> ndarray: ...
+) -> NDArray[Any]: ...
 @overload
 def linspace(
-    start: _ArrayLikeNumber,
-    stop: _ArrayLikeNumber,
+    start: _ArrayLikeFloat_co,
+    stop: _ArrayLikeFloat_co,
     num: SupportsIndex = ...,
     endpoint: bool = ...,
-    retstep: Literal[True] = ...,
+    retstep: L[True] = ...,
+    dtype: None = ...,
+    axis: SupportsIndex = ...,
+) -> Tuple[NDArray[floating[Any]], floating[Any]]: ...
+@overload
+def linspace(
+    start: _ArrayLikeComplex_co,
+    stop: _ArrayLikeComplex_co,
+    num: SupportsIndex = ...,
+    endpoint: bool = ...,
+    retstep: L[True] = ...,
+    dtype: None = ...,
+    axis: SupportsIndex = ...,
+) -> Tuple[NDArray[complexfloating[Any, Any]], complexfloating[Any, Any]]: ...
+@overload
+def linspace(
+    start: _ArrayLikeComplex_co,
+    stop: _ArrayLikeComplex_co,
+    num: SupportsIndex = ...,
+    endpoint: bool = ...,
+    retstep: L[True] = ...,
+    dtype: _DTypeLike[_SCT] = ...,
+    axis: SupportsIndex = ...,
+) -> Tuple[NDArray[_SCT], _SCT]: ...
+@overload
+def linspace(
+    start: _ArrayLikeComplex_co,
+    stop: _ArrayLikeComplex_co,
+    num: SupportsIndex = ...,
+    endpoint: bool = ...,
+    retstep: L[True] = ...,
     dtype: DTypeLike = ...,
     axis: SupportsIndex = ...,
-) -> Tuple[ndarray, Any]: ...
+) -> Tuple[NDArray[Any], Any]: ...
 
+@overload
 def logspace(
-    start: _ArrayLikeNumber,
-    stop: _ArrayLikeNumber,
+    start: _ArrayLikeFloat_co,
+    stop: _ArrayLikeFloat_co,
     num: SupportsIndex = ...,
     endpoint: bool = ...,
-    base: _ArrayLikeNumber = ...,
+    base: _ArrayLikeFloat_co = ...,
+    dtype: None = ...,
+    axis: SupportsIndex = ...,
+) -> NDArray[floating[Any]]: ...
+@overload
+def logspace(
+    start: _ArrayLikeComplex_co,
+    stop: _ArrayLikeComplex_co,
+    num: SupportsIndex = ...,
+    endpoint: bool = ...,
+    base: _ArrayLikeComplex_co = ...,
+    dtype: None = ...,
+    axis: SupportsIndex = ...,
+) -> NDArray[complexfloating[Any, Any]]: ...
+@overload
+def logspace(
+    start: _ArrayLikeComplex_co,
+    stop: _ArrayLikeComplex_co,
+    num: SupportsIndex = ...,
+    endpoint: bool = ...,
+    base: _ArrayLikeComplex_co = ...,
+    dtype: _DTypeLike[_SCT] = ...,
+    axis: SupportsIndex = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def logspace(
+    start: _ArrayLikeComplex_co,
+    stop: _ArrayLikeComplex_co,
+    num: SupportsIndex = ...,
+    endpoint: bool = ...,
+    base: _ArrayLikeComplex_co = ...,
     dtype: DTypeLike = ...,
     axis: SupportsIndex = ...,
-) -> ndarray: ...
+) -> NDArray[Any]: ...
 
+@overload
+def geomspace(
+    start: _ArrayLikeFloat_co,
+    stop: _ArrayLikeFloat_co,
+    num: SupportsIndex = ...,
+    endpoint: bool = ...,
+    dtype: None = ...,
+    axis: SupportsIndex = ...,
+) -> NDArray[floating[Any]]: ...
+@overload
+def geomspace(
+    start: _ArrayLikeComplex_co,
+    stop: _ArrayLikeComplex_co,
+    num: SupportsIndex = ...,
+    endpoint: bool = ...,
+    dtype: None = ...,
+    axis: SupportsIndex = ...,
+) -> NDArray[complexfloating[Any, Any]]: ...
+@overload
+def geomspace(
+    start: _ArrayLikeComplex_co,
+    stop: _ArrayLikeComplex_co,
+    num: SupportsIndex = ...,
+    endpoint: bool = ...,
+    dtype: _DTypeLike[_SCT] = ...,
+    axis: SupportsIndex = ...,
+) -> NDArray[_SCT]: ...
+@overload
 def geomspace(
-    start: _ArrayLikeNumber,
-    stop: _ArrayLikeNumber,
+    start: _ArrayLikeComplex_co,
+    stop: _ArrayLikeComplex_co,
     num: SupportsIndex = ...,
     endpoint: bool = ...,
     dtype: DTypeLike = ...,
     axis: SupportsIndex = ...,
-) -> ndarray: ...
+) -> NDArray[Any]: ...
 
 # Re-exported to `np.lib.function_base`
 def add_newdoc(