summaryrefslogtreecommitdiff
path: root/numpy
diff options
context:
space:
mode:
Diffstat (limited to 'numpy')
-rw-r--r--numpy/core/function_base.pyi184
-rw-r--r--numpy/typing/tests/data/fail/array_constructors.pyi6
-rw-r--r--numpy/typing/tests/data/reveal/array_constructors.pyi22
3 files changed, 183 insertions, 29 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(
diff --git a/numpy/typing/tests/data/fail/array_constructors.pyi b/numpy/typing/tests/data/fail/array_constructors.pyi
index 4f0a60b5b..065b7d8a0 100644
--- a/numpy/typing/tests/data/fail/array_constructors.pyi
+++ b/numpy/typing/tests/data/fail/array_constructors.pyi
@@ -21,10 +21,10 @@ np.linspace(0, 2, retstep=b'False') # E: No overload variant
np.linspace(0, 2, dtype=0) # E: No overload variant
np.linspace(0, 2, axis=None) # E: No overload variant
-np.logspace(None, 'bob') # E: Argument 1
-np.logspace(0, 2, base=None) # E: Argument "base"
+np.logspace(None, 'bob') # E: No overload variant
+np.logspace(0, 2, base=None) # E: No overload variant
-np.geomspace(None, 'bob') # E: Argument 1
+np.geomspace(None, 'bob') # E: No overload variant
np.stack(generator) # E: No overload variant
np.hstack({1, 2}) # E: No overload variant
diff --git a/numpy/typing/tests/data/reveal/array_constructors.pyi b/numpy/typing/tests/data/reveal/array_constructors.pyi
index 233988e63..ba5710e0f 100644
--- a/numpy/typing/tests/data/reveal/array_constructors.pyi
+++ b/numpy/typing/tests/data/reveal/array_constructors.pyi
@@ -114,10 +114,24 @@ reveal_type(np.require(B, requirements="W")) # E: SubClass[{float64}]
reveal_type(np.require(B, requirements="A")) # E: SubClass[{float64}]
reveal_type(np.require(C)) # E: ndarray[Any, Any]
-reveal_type(np.linspace(0, 10)) # E: ndarray[Any, Any]
-reveal_type(np.linspace(0, 10, retstep=True)) # E: Tuple[ndarray[Any, Any], Any]
-reveal_type(np.logspace(0, 10)) # E: ndarray[Any, Any]
-reveal_type(np.geomspace(1, 10)) # E: ndarray[Any, Any]
+reveal_type(np.linspace(0, 10)) # E: ndarray[Any, dtype[floating[Any]]]
+reveal_type(np.linspace(0, 10j)) # E: ndarray[Any, dtype[complexfloating[Any, Any]]]
+reveal_type(np.linspace(0, 10, dtype=np.int64)) # E: ndarray[Any, dtype[{int64}]]
+reveal_type(np.linspace(0, 10, dtype=int)) # E: ndarray[Any, dtype[Any]]
+reveal_type(np.linspace(0, 10, retstep=True)) # E: Tuple[ndarray[Any, dtype[floating[Any]]], floating[Any]]
+reveal_type(np.linspace(0j, 10, retstep=True)) # E: Tuple[ndarray[Any, dtype[complexfloating[Any, Any]]], complexfloating[Any, Any]]
+reveal_type(np.linspace(0, 10, retstep=True, dtype=np.int64)) # E: Tuple[ndarray[Any, dtype[{int64}]], {int64}]
+reveal_type(np.linspace(0j, 10, retstep=True, dtype=int)) # E: Tuple[ndarray[Any, dtype[Any]], Any]
+
+reveal_type(np.logspace(0, 10)) # E: ndarray[Any, dtype[floating[Any]]]
+reveal_type(np.logspace(0, 10j)) # E: ndarray[Any, dtype[complexfloating[Any, Any]]]
+reveal_type(np.logspace(0, 10, dtype=np.int64)) # E: ndarray[Any, dtype[{int64}]]
+reveal_type(np.logspace(0, 10, dtype=int)) # E: ndarray[Any, dtype[Any]]
+
+reveal_type(np.geomspace(0, 10)) # E: ndarray[Any, dtype[floating[Any]]]
+reveal_type(np.geomspace(0, 10j)) # E: ndarray[Any, dtype[complexfloating[Any, Any]]]
+reveal_type(np.geomspace(0, 10, dtype=np.int64)) # E: ndarray[Any, dtype[{int64}]]
+reveal_type(np.geomspace(0, 10, dtype=int)) # E: ndarray[Any, dtype[Any]]
reveal_type(np.zeros_like(A)) # E: ndarray[Any, dtype[{float64}]]
reveal_type(np.zeros_like(C)) # E: ndarray[Any, dtype[Any]]
View Lorry Controller Status