8000 🐛 fix typing errors in `_core.shape_base` (#108) · numpy/numtype@ae7b41a · GitHub
[go: up one dir, main page]
Skip to content

Commit ae7b41a

Browse files
jorenhamjorenham
authored
🐛 fix typing errors in _core.shape_base (#108)
1 parent 28bbef2 commit ae7b41a

File tree

2 files changed

+176
-74
lines changed

2 files changed

+176
-74
lines changed

‎src/numpy-stubs/_core/shape_base.pyi

Lines changed: 132 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,98 +1,196 @@
11
from collections.abc import Sequence
2-
from typing import Any, SupportsIndex, overload
3-
from typing_extensions import TypeVar
2+
from typing import Any, SupportsIndex, TypeAlias, overload
3+
from typing_extensions import TypeVar, Unpack
44

55
import numpy as np
66
from numpy._typing import ArrayLike, DTypeLike, NDArray, _ArrayLike, _DTypeLike
77

88
__all__ = ["atleast_1d", "atleast_2d", "atleast_3d", "block", "hstack", "stack", "unstack", "vstack"]
99

10+
###
11+
1012
_SCT = TypeVar("_SCT", bound=np.generic)
13+
_SCT0 = TypeVar("_SCT0", bound=np.generic)
14+
_SCT1 = TypeVar("_SCT1", bound=np.generic)
15+
1116
_ArrayT = TypeVar("_ArrayT", bound=NDArray[Any])
1217

18+
_Array1T = TypeVar("_Array1T", bound=np.ndarray[_AtLeast1D, np.dtype[Any]])
19+
_Array1T0 = TypeVar("_Array1T0", bound=np.ndarray[_AtLeast1D, np.dtype[Any]])
20+
_Array1T1 = TypeVar("_Array1T1", bound=np.ndarray[_AtLeast1D, np.dtype[Any]])
21+
22+
_Array2T = TypeVar("_Array2T", bound=np.ndarray[_AtLeast2D, np.dtype[Any]])
23+
_Array2T0 = TypeVar("_Array2T0", bound=np.ndarray[_AtLeast2D, np.dtype[Any]])
24+
_Array2T1 = TypeVar("_Array2T1", bound=np.ndarray[_AtLeast2D, np.dtype[Any]])
25+
26+
_Array3T = TypeVar("_Array3T", bound=np.ndarray[_AtLeast3D, np.dtype[Any]])
27+
_Array3T0 = TypeVar("_Array3T0", bound=np.ndarray[_AtLeast3D, np.dtype[Any]])
28+
_Array3T1 = TypeVar("_Array3T1", bound=np.ndarray[_AtLeast3D, np.dtype[Any]])
29+
30+
_AtLeast1D: TypeAlias = tuple[int, Unpack[tuple[int, ...]]]
31+
_AtLeast2D: TypeAlias = tuple[int, int, Unpack[tuple[int, ...]]]
32+
_AtLeast3D: TypeAlias = tuple[int, int, int, Unpack[tuple[int, ...]]]
33+
34+
###
35+
1336
#
1437
@overload
15-
def atleast_1d(arys: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
38+
def atleast_1d(a0: _Array1T, /) -> _Array1T: ...
39+
@overload
40+
def atleast_1d(a0: _Array1T0, a1: _Array1T1, /) -> tuple[_Array1T0, _Array1T1]: ...
41+
@overload
42+
def atleast_1d(a0: _Array1T, a1: _Array1T, /, *arys: _Array1T) -> tuple[_Array1T, ...]: ...
43+
@overload
44+
def atleast_1d(a0: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
45+
@overload
46+
def atleast_1d(a0: _ArrayLike[_SCT0], a2: _ArrayLike[_SCT1], /) -> tuple[NDArray[_SCT0], NDArray[_SCT1]]: ...
1647
@overload
17-
def atleast_1d(arys: ArrayLike, /) -> NDArray[Any]: ...
48+
def atleast_1d(a0: _ArrayLike[_SCT], a2: _ArrayLike[_SCT], /, *arys: _ArrayLike[_SCT]) -> tuple[NDArray[_SCT], ...]: ...
1849
@overload
19-
def atleast_1d(*arys: ArrayLike) -> tuple[NDArray[Any], ...]: ...
50+
def atleast_1d(a0: ArrayLike, /) -> NDArray[Any]: ...
51+
@overload
52+
def atleast_1d(a0: ArrayLike, a2: ArrayLike, /) -> tuple[NDArray[Any], NDArray[Any]]: ...
53+
@overload
54+
def atleast_1d(a0: ArrayLike, a2: ArrayLike, /, *arys: ArrayLike) -> tuple[NDArray[Any], ...]: ...
2055

2156
#
2257
@overload
23-
def atleast_2d(arys: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
58+
def atleast_2d(a0: _Array2T, /) -> _Array2T: ...
59+
@overload
60+
def atleast_2d(a0: _Array2T0, a1: _Array2T1, /) -> tuple[_Array2T0, _Array2T1]: ...
61+
@overload
62+
def atleast_2d(a0: _Array2T, a1: _Array2T, /, *arys: _Array2T) -> tuple[_Array2T, ...]: ...
63+
@overload
64+
def atleast_2d(a0: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
65+
@overload
66+
def atleast_2d(a0: _ArrayLike[_SCT0], a2: _ArrayLike[_SCT1], /) -> tuple[NDArray[_SCT0], NDArray[_SCT1]]: ...
2467
@overload
25-
def atleast_2d(arys: ArrayLike, /) -> NDArray[Any]: ...
68+
def atleast_2d(a0: _ArrayLike[_SCT], a2: _ArrayLike[_SCT], /, *arys: _ArrayLike[_SCT]) -> tuple[NDArray[_SCT], ...]: ...
2669
@overload
27-
def atleast_2d(*arys: ArrayLike) -> tuple[NDArray[Any], ...]: ...
70+
def atleast_2d(a0: ArrayLike, /) -> NDArray[Any]: ...
71+
@overload
72+
def atleast_2d(a0: ArrayLike, a2: ArrayLike, /) -> tuple[NDArray[Any], NDArray[Any]]: ...
73+
@overload
74+
def atleast_2d(a0: ArrayLike, a2: ArrayLike, /, *arys: ArrayLike) -> tuple[NDArray[Any], ...]: ...
2875

2976
#
3077
@overload
31-
def atleast_3d(arys: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
78+
def atleast_3d(a0: _Array3T, /) -> _Array3T: ...
79+
@overload
80+
def atleast_3d(a0: _Array3T0, a1: _Array3T1, /) -> tuple[_Array3T0, _Array3T1]: ...
81+
@overload
82+
def atleast_3d(a0: _Array3T, a1: _Array3T, /, *arys: _Array3T) -> tuple[_Array3T, ...]: ...
83+
@overload
84+
def atleast_3d(a0: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
3285
@overload
33-
def atleast_3d(arys: ArrayLike, /) -> NDArray[Any]: ...
86+
def atleast_3d(a0: _ArrayLike[_SCT0], a2: _ArrayLike[_SCT1], /) -> tuple[NDArray[_SCT0], NDArray[_SCT1]]: ...
3487
@overload
35-
def atleast_3d(*arys: ArrayLike) -> tuple[NDArray[Any], ...]: ...
88+
def atleast_3d(a0: _ArrayLike[_SCT], a2: _ArrayLike[_SCT], /, *arys: _ArrayLike[_SCT]) -> tuple[NDArray[_SCT], ...]: ...
89+
@overload
90+
def atleast_3d(a0: ArrayLike, /) -> NDArray[Any]: ...
91+
@overload
92+
def atleast_3d(a0: ArrayLike, a2: ArrayLike, /) -> tuple[NDArray[Any], NDArray[Any]]: ...
93+
@overload
94+
def atleast_3d(a0: ArrayLike, a2: ArrayLike, /, *arys: ArrayLike) -> tuple[NDArray[Any], ...]: ...
3695

3796
#
3897
@overload
39-
def vstack(tup: Sequence[_ArrayLike[_SCT]], *, dtype: None = ..., casting: np._CastingKind = ...) -> NDArray[_SCT]: ...
98+
def vstack(
99+
tup: Sequence[_ArrayLike[_SCT]],
100+
*,
101+
dtype: None = None,
102+
casting: np._CastingKind = "same_kind",
103+
) -> NDArray[_SCT]: ...
40104
@overload
41-
def vstack(tup: Sequence[ArrayLike], *, dtype: _DTypeLike[_SCT], casting: np._CastingKind = ...) -> NDArray[_SCT]: ...
105+
def vstack(
106+
tup: Sequence[ArrayLike],
107+
*,
108+
dtype: _DTypeLike[_SCT],
109+
casting: np._CastingKind = "same_kind",
110+
) -> NDArray[_SCT]: ...
42111
@overload
43-
def vstack(tup: Sequence[ArrayLike], *, dtype: DTypeLike = ..., casting: np._CastingKind = ...) -> NDArray[Any]: ...
112+
def vstack(
113+
tup: Sequence[ArrayLike],
114+
*,
115+
dtype: DTypeLike | None = None,
116+
casting: np._CastingKind = "same_kind",
117+
) -> NDArray[Any]: ...
44118

45119
#
46120
@overload
47-
def hstack(tup: Sequence[_ArrayLike[_SCT]], *, dtype: None = ..., casting: np._CastingKind = ...) -> NDArray[_SCT]: ...
121+
def hstack(
122+
tup: Sequence[_ArrayLike[_SCT]],
123+
*,
124+
dtype: None = None,
125+
casting: np._CastingKind = "same_kind",
126+
) -> NDArray[_SCT]: ...
48127
@overload
49-
def hstack(tup: Sequence[ArrayLike], *, dtype: _DTypeLike[_SCT], casting: np._CastingKind = ...) -> NDArray[_SCT]: ...
128+
def hstack(
129+
tup: Sequence[ArrayLike],
130+
*,
131+
dtype: _DTypeLike[_SCT],
132+
casting: np._CastingKind = "same_kind",
133+
) -> NDArray[_SCT]: ...
50134
@overload
51-
def hstack(tup: Sequence[ArrayLike], *, dtype: DTypeLike = ..., casting: np._CastingKind = ...) -> NDArray[Any]: ...
135+
def hstack(
136+
tup: Sequence[ArrayLike],
137+
*,
138+
dtype: DTypeLike | None = None,
139+
casting: np._CastingKind = "same_kind",
140+
) -> NDArray[Any]: ...
52141

53142
#
54143
@overload
55144
def stack(
56145
arrays: Sequence[_ArrayLike[_SCT]],
57-
axis: SupportsIndex = ...,
58-
out: None = ...,
146+
axis: SupportsIndex = 0,
147+
out: None = None,
59148
*,
60-
dtype: None = ...,
61-
casting: np._CastingKind = ...,
149+
dtype: None = None,
150+
casting: np._CastingKind = "same_kind",
62151
) -> NDArray[_SCT]: ...
63152
@overload
64153
def stack(
65154
arrays: Sequence[ArrayLike],
66-
axis: SupportsIndex = ...,
67-
out: None = ...,
155+
axis: SupportsIndex = 0,
156+
out: None = None,
68157
*,
69158
dtype: _DTypeLike[_SCT],
70-
casting: np._CastingKind = ...,
159+
casting: np._CastingKind = "same_kind",
71160
) -> NDArray[_SCT]: ...
72161
@overload
73162
def stack(
74163
arrays: Sequence[ArrayLike],
75-
axis: SupportsIndex = ...,
76-
out: None = ...,
164+
axis: SupportsIndex = 0,
165+
out: None = None,
77166
*,
78-
dtype: DTypeLike = ...,
79-
casting: np._CastingKind = ...,
167+
dtype: DTypeLike | None = None,
168+
casting: np._CastingKind = "same_kind",
80169
) -> NDArray[Any]: ...
81170
@overload
82171
def stack(
83172
arrays: Sequence[ArrayLike],
84-
axis: SupportsIndex = ...,
85-
out: _ArrayT = ...,
173+
axis: SupportsIndex,
174+
out: _ArrayT,
175+
*,
176+
dtype: DTypeLike | None = None,
177+
casting: np._CastingKind = "same_kind",
178+
) -> _ArrayT: ...
179+
@overload
180+
def stack(
181+
arrays: Sequence[ArrayLike],
182+
axis: SupportsIndex = 0,
86183
*,
87-
dtype: DTypeLike = ...,
88-
casting: np._CastingKind = ...,
184+
out: _ArrayT,
185+
dtype: DTypeLike | None = None,
186+
casting: np._CastingKind = "same_kind",
89187
) -> _ArrayT: ...
90188

91189
#
92190
@overload
93-
def unstack(array: _ArrayLike[_SCT], /, *, axis: int = ...) -> tuple[NDArray[_SCT], ...]: ...
191+
def unstack(array: _ArrayLike[_SCT], /, *, axis: SupportsIndex = 0) -> tuple[NDArray[_SCT], ...]: ...
94192
@overload
95-
def unstack(array: ArrayLike, /, *, axis: int = ...) -> tuple[NDArray[Any], ...]: ...
193+
def unstack(array: ArrayLike, /, *, axis: SupportsIndex = 0) -> tuple[NDArray[Any], ...]: ...
96194

97195
#
98196
@overload

0 commit comments

Comments
 (0)
0