8000 Merge pull request #19245 from BvB93/shape_base · numpy/numpy@e0ceabc · GitHub
[go: up one dir, main page]
Skip to content

Commit e0ceabc

Browse files
charrischarris
authored
Merge pull request #19245 from BvB93/shape_base
ENH: Add dtype-support to the `np.core.shape_base` annotations
2 parents a914cb7 + 76bad56 commit e0ceabc

File tree

3 files changed

+69
-34
lines changed

3 files changed

+69
-34
lines changed

‎numpy/core/shape_base.pyi

Lines changed: 49 additions & 16 deletions
< 57AE tr class="diff-line-row">
Original file line numberDiff line numberDiff line change
@@ -1,39 +1,72 @@
11
import sys
2-
from typing import TypeVar, overload, List, Sequence
2+
from typing import TypeVar, overload, List, Sequence, Any
33

4-
from numpy import ndarray
5-
from numpy.typing import ArrayLike
4+
from numpy import generic, dtype
5+
from numpy.typing import ArrayLike, NDArray, _NestedSequence, _SupportsArray
66

77
if sys.version_info >= (3, 8):
88
from typing import SupportsIndex
99
else:
1010
from typing_extensions import SupportsIndex
1111

12-
_ArrayType = TypeVar("_ArrayType", bound=ndarray)
12+
_SCT = TypeVar("_SCT", bound=generic)
13+
_ArrayType = TypeVar("_ArrayType", bound=NDArray[Any])
1314

15+
_ArrayLike = _NestedSequence[_SupportsArray[dtype[_SCT]]]
16+
17+
__all__: List[str]
18+
19+
@overload
20+
def atleast_1d(__arys: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...
21+
@overload
22+
def atleast_1d(__arys: ArrayLike) -> NDArray[Any]: ...
23+
@overload
24+
def atleast_1d(*arys: ArrayLike) -> List[NDArray[Any]]: ...
25+
26+
@overload
27+
def atleast_2d(__arys: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...
28+
@overload
29+
def atleast_2d(__arys: ArrayLike) -> NDArray[Any]: ...
1430
@overload
15-
def atleast_1d(__arys: ArrayLike) -> ndarray: ...
31+
def atleast_2d(*arys: ArrayLike) -> List[NDArray[Any]]: ...
32+
33+
@overload
34+
def atleast_3d(__arys: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...
1635
@overload
17-
def atleast_1d(*arys: ArrayLike) -> List[ndarray]: ...
36+
def atleast_3d(__arys: ArrayLike) -> NDArray[Any]: ...
37+
@overload
38+
def atleast_3d(*arys: ArrayLike) -> List[NDArray[Any]]: ...
1839

1940
@overload
20-
def atleast_2d(__arys: ArrayLike) -> ndarray: ...
41+
def vstack(tup: Sequence[_ArrayLike[_SCT]]) -> NDArray[_SCT]: ...
2142
@overload
22-
def atleast_2d(*arys: ArrayLike) -> List[ndarray]: ...
43+
def vstack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ...
2344

2445
@overload
25-
def atleast_3d(__arys: ArrayLike) -> ndarray: ...
46+
def hstack(tup: Sequence[_ArrayLike[_SCT]]) -> NDArray[_SCT]: ...
2647
@overload
27-
def atleast_3d(*arys: ArrayLike) -> List[ndarray]: ...
48+
def hstack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ...
2849

29-
def vstack(tup: Sequence[ArrayLike]) -> ndarray: ...
30-
def hstack(tup: Sequence[ArrayLike]) -> ndarray: ...
3150
@overload
3251
def stack(
33-
arrays: Sequence[ArrayLike], axis: SupportsIndex = ..., out: None = ...
34-
) -> ndarray: ...
52+
arrays: Sequence[_ArrayLike[_SCT]],
53+
axis: SupportsIndex = ...,
54+
out: None = ...,
55+
) -> NDArray[_SCT]: ...
56+
@overload
57+
def stack(
58+
arrays: Sequence[ArrayLike],
59+
axis: SupportsIndex = ...,
60+
out: None = ...,
61+
) -> NDArray[Any]: ...
3562
@overload
3663
def stack(
37-
arrays: Sequence[ArrayLike], axis: SupportsIndex = ..., out: _ArrayType = ...
64+
arrays: Sequence[ArrayLike],
65+
axis: SupportsIndex = ...,
66+
out: _ArrayType = ...,
3867
) -> _ArrayType: ...
39-
def block(arrays: ArrayLike) -> ndarray: ...
68+
69+
@overload
70+
def block(arrays: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...
71+
@overload
72+
def block(arrays: ArrayLike) -> NDArray[Any]: ...

‎numpy/typing/tests/data/fail/array_constructors.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,5 +27,5 @@
2727
np.geomspace(None, 'bob') # E: Argument 1
2828

2929
np.stack(generator) # E: No overload variant
30-
np.hstack({1, 2}) # E: incompatible type
31-
np.vstack(1) # E: incompatible type
30+
np.hstack({1, 2}) # E: No overload variant
31+
np.vstack(1) # E: No overload variant

‎numpy/typing/tests/data/reveal/array_constructors.py

Lines changed: 18 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -147,25 +147,27 @@ def func(i: int, j: int, **kwargs: Any) -> SubClass[np.float64]: ...
147147

148148
reveal_type(np.identity(10)) # E: numpy.ndarray[Any, Any]
149149

150-
reveal_type(np.atleast_1d(A)) # E: numpy.ndarray[Any, Any]
151-
reveal_type(np.atleast_1d(C)) # E: numpy.ndarray[Any, Any]
152-
reveal_type(np.atleast_1d(A, A)) # E: list[numpy.ndarray[Any, Any]]
153-
reveal_type(np.atleast_1d(A, C)) # E: list[numpy.ndarray[Any, Any]]
154-
reveal_type(np.atleast_1d(C, C)) # E: list[numpy.ndarray[Any, Any]]
150+
reveal_type(np.atleast_1d(A)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
151+
reveal_type(np.atleast_1d(C)) # E: numpy.ndarray[Any, numpy.dtype[Any]]
152+
reveal_type(np.atleast_1d(A, A)) # E: list[numpy.ndarray[Any, numpy.dtype[Any]]]
153+
reveal_type(np.atleast_1d(A, C)) # E: list[numpy.ndarray[Any, numpy.dtype[Any]]]
154+
reveal_type(np.atleast_1d(C, C)) # E: list[numpy.ndarray[Any, numpy.dtype[Any]]]
155155

156-
reveal_type(np.atleast_2d(A)) # E: numpy.ndarray[Any, Any]
156+
reveal_type(np.atleast_2d(A)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
157157

158-
reveal_type(np.atleast_3d(A)) # E: numpy.ndarray[Any, Any]
158+
reveal_type(np.atleast_3d(A)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
159159

160-
reveal_type(np.vstack([A, A])) # E: numpy.ndarray[Any, Any]
161-
reveal_type(np.vstack([A, C])) # E: numpy.ndarray[Any, Any]
162-
reveal_type(np.vstack([C, C])) # E: numpy.ndarray[Any, Any]
160+
reveal_type(np.vstack([A, A])) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
161+
reveal_type(np.vstack([A, C])) # E: numpy.ndarray[Any, numpy.dtype[Any]]
162+
reveal_type(np.vstack([C, C])) # E: numpy.ndarray[Any, numpy.dtype[Any]]
163163

164-
reveal_type(np.hstack([A, A])) # E: numpy.ndarray[Any, Any]
164+
reveal_type(np.hstack([A, A])) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
165165

166-
reveal_type(np.stack([A, A])) # E: numpy.ndarray[Any, Any]
167-
reveal_type(np.stack([A, A], axis=0)) # E: numpy.ndarray[Any, Any]
168-
reveal_type(np.stack([A, A], out=B)) # E: SubClass
166+
reveal_type(np.stack([A, A])) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
167+
reveal_type(np.stack([A, C])) # E: numpy.ndarray[Any, numpy.dtype[Any]]
168+
reveal_type(np.stack([C, C])) # E: numpy.ndarray[Any, numpy.dtype[Any]]
169+
reveal_type(np.stack([A, A], axis=0)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
170+
reveal_type(np.stack([A, A], out=B)) # E: SubClass[{float64}]
169171

170-
reveal_type(np.block([[A, A], [A, A]])) # E: numpy.ndarray[Any, Any]
171-
reveal_type(np.block(C)) # E: numpy.ndarray[Any, Any]
172+
reveal_type(np.block([[A, A], [A, A]])) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
173+
reveal_type(np.block(C)) # E: numpy.ndarray[Any, numpy.dtype[Any]]

0 commit comments

Comments
 (0)
0