diff options
| -rw-r--r-- | numpy/core/shape_base.pyi | 65 | ||||
| -rw-r--r-- | numpy/typing/tests/data/fail/array_constructors.py | 4 | ||||
| -rw-r--r-- | numpy/typing/tests/data/reveal/array_constructors.py | 34 |
3 files changed, 69 insertions, 34 deletions
diff --git a/numpy/core/shape_base.pyi b/numpy/core/shape_base.pyi index ec40a8814..9aaeceed7 100644 --- a/numpy/core/shape_base.pyi +++ b/numpy/core/shape_base.pyi @@ -1,39 +1,72 @@ import sys -from typing import TypeVar, overload, List, Sequence +from typing import TypeVar, overload, List, Sequence, Any -from numpy import ndarray -from numpy.typing import ArrayLike +from numpy import generic, dtype +from numpy.typing import ArrayLike, NDArray, _NestedSequence, _SupportsArray if sys.version_info >= (3, 8): from typing import SupportsIndex else: from typing_extensions import SupportsIndex -_ArrayType = TypeVar("_ArrayType", bound=ndarray) +_SCT = TypeVar("_SCT", bound=generic) +_ArrayType = TypeVar("_ArrayType", bound=NDArray[Any]) +_ArrayLike = _NestedSequence[_SupportsArray[dtype[_SCT]]] + +__all__: List[str] + +@overload +def atleast_1d(__arys: _ArrayLike[_SCT]) -> NDArray[_SCT]: ... +@overload +def atleast_1d(__arys: ArrayLike) -> NDArray[Any]: ... +@overload +def atleast_1d(*arys: ArrayLike) -> List[NDArray[Any]]: ... + +@overload +def atleast_2d(__arys: _ArrayLike[_SCT]) -> NDArray[_SCT]: ... +@overload +def atleast_2d(__arys: ArrayLike) -> NDArray[Any]: ... @overload -def atleast_1d(__arys: ArrayLike) -> ndarray: ... +def atleast_2d(*arys: ArrayLike) -> List[NDArray[Any]]: ... + +@overload +def atleast_3d(__arys: _ArrayLike[_SCT]) -> NDArray[_SCT]: ... @overload -def atleast_1d(*arys: ArrayLike) -> List[ndarray]: ... +def atleast_3d(__arys: ArrayLike) -> NDArray[Any]: ... +@overload +def atleast_3d(*arys: ArrayLike) -> List[NDArray[Any]]: ... @overload -def atleast_2d(__arys: ArrayLike) -> ndarray: ... +def vstack(tup: Sequence[_ArrayLike[_SCT]]) -> NDArray[_SCT]: ... @overload -def atleast_2d(*arys: ArrayLike) -> List[ndarray]: ... +def vstack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ... @overload -def atleast_3d(__arys: ArrayLike) -> ndarray: ... +def hstack(tup: Sequence[_ArrayLike[_SCT]]) -> NDArray[_SCT]: ... @overload -def atleast_3d(*arys: ArrayLike) -> List[ndarray]: ... +def hstack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ... -def vstack(tup: Sequence[ArrayLike]) -> ndarray: ... -def hstack(tup: Sequence[ArrayLike]) -> ndarray: ... @overload def stack( - arrays: Sequence[ArrayLike], axis: SupportsIndex = ..., out: None = ... -) -> ndarray: ... + arrays: Sequence[_ArrayLike[_SCT]], + axis: SupportsIndex = ..., + out: None = ..., +) -> NDArray[_SCT]: ... +@overload +def stack( + arrays: Sequence[ArrayLike], + axis: SupportsIndex = ..., + out: None = ..., +) -> NDArray[Any]: ... @overload def stack( - arrays: Sequence[ArrayLike], axis: SupportsIndex = ..., out: _ArrayType = ... + arrays: Sequence[ArrayLike], + axis: SupportsIndex = ..., + out: _ArrayType = ..., ) -> _ArrayType: ... -def block(arrays: ArrayLike) -> ndarray: ... + +@overload +def block(arrays: _ArrayLike[_SCT]) -> NDArray[_SCT]: ... +@overload +def block(arrays: ArrayLike) -> NDArray[Any]: ... diff --git a/numpy/typing/tests/data/fail/array_constructors.py b/numpy/typing/tests/data/fail/array_constructors.py index eb57e5c00..0e2250513 100644 --- a/numpy/typing/tests/data/fail/array_constructors.py +++ b/numpy/typing/tests/data/fail/array_constructors.py @@ -27,5 +27,5 @@ np.logspace(0, 2, base=None) # E: Argument "base" np.geomspace(None, 'bob') # E: Argument 1 np.stack(generator) # E: No overload variant -np.hstack({1, 2}) # E: incompatible type -np.vstack(1) # E: incompatible type +np.hstack({1, 2}) # E: No overload variant +np.vstack(1) # E: No overload variant diff --git a/numpy/typing/tests/data/reveal/array_constructors.py b/numpy/typing/tests/data/reveal/array_constructors.py index 1b9006220..44c85e988 100644 --- a/numpy/typing/tests/data/reveal/array_constructors.py +++ b/numpy/typing/tests/data/reveal/array_constructors.py @@ -147,25 +147,27 @@ reveal_type(np.fromfunction(func, (3, 5))) # E: SubClass[{float64}] reveal_type(np.identity(10)) # E: numpy.ndarray[Any, Any] -reveal_type(np.atleast_1d(A)) # E: numpy.ndarray[Any, Any] -reveal_type(np.atleast_1d(C)) # E: numpy.ndarray[Any, Any] -reveal_type(np.atleast_1d(A, A)) # E: list[numpy.ndarray[Any, Any]] -reveal_type(np.atleast_1d(A, C)) # E: list[numpy.ndarray[Any, Any]] -reveal_type(np.atleast_1d(C, C)) # E: list[numpy.ndarray[Any, Any]] +reveal_type(np.atleast_1d(A)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]] +reveal_type(np.atleast_1d(C)) # E: numpy.ndarray[Any, numpy.dtype[Any]] +reveal_type(np.atleast_1d(A, A)) # E: list[numpy.ndarray[Any, numpy.dtype[Any]]] +reveal_type(np.atleast_1d(A, C)) # E: list[numpy.ndarray[Any, numpy.dtype[Any]]] +reveal_type(np.atleast_1d(C, C)) # E: list[numpy.ndarray[Any, numpy.dtype[Any]]] -reveal_type(np.atleast_2d(A)) # E: numpy.ndarray[Any, Any] +reveal_type(np.atleast_2d(A)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]] -reveal_type(np.atleast_3d(A)) # E: numpy.ndarray[Any, Any] +reveal_type(np.atleast_3d(A)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]] -reveal_type(np.vstack([A, A])) # E: numpy.ndarray[Any, Any] -reveal_type(np.vstack([A, C])) # E: numpy.ndarray[Any, Any] -reveal_type(np.vstack([C, C])) # E: numpy.ndarray[Any, Any] +reveal_type(np.vstack([A, A])) # E: numpy.ndarray[Any, numpy.dtype[{float64}]] +reveal_type(np.vstack([A, C])) # E: numpy.ndarray[Any, numpy.dtype[Any]] +reveal_type(np.vstack([C, C])) # E: numpy.ndarray[Any, numpy.dtype[Any]] -reveal_type(np.hstack([A, A])) # E: numpy.ndarray[Any, Any] +reveal_type(np.hstack([A, A])) # E: numpy.ndarray[Any, numpy.dtype[{float64}]] -reveal_type(np.stack([A, A])) # E: numpy.ndarray[Any, Any] -reveal_type(np.stack([A, A], axis=0)) # E: numpy.ndarray[Any, Any] -reveal_type(np.stack([A, A], out=B)) # E: SubClass +reveal_type(np.stack([A, A])) # E: numpy.ndarray[Any, numpy.dtype[{float64}]] +reveal_type(np.stack([A, C])) # E: numpy.ndarray[Any, numpy.dtype[Any]] +reveal_type(np.stack([C, C])) # E: numpy.ndarray[Any, numpy.dtype[Any]] +reveal_type(np.stack([A, A], axis=0)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]] +reveal_type(np.stack([A, A], out=B)) # E: SubClass[{float64}] -reveal_type(np.block([[A, A], [A, A]])) # E: numpy.ndarray[Any, Any] -reveal_type(np.block(C)) # E: numpy.ndarray[Any, Any] +reveal_type(np.block([[A, A], [A, A]])) # E: numpy.ndarray[Any, numpy.dtype[{float64}]] +reveal_type(np.block(C)) # E: numpy.ndarray[Any, numpy.dtype[Any]] |