codekingpro/portable-devtools

1from collections.abc import Sequence
2from typing import Any, SupportsIndex, TypeVar, overload
3
4from numpy import _CastingKind, generic
5from numpy._typing import ArrayLike, DTypeLike, NDArray, _ArrayLike, _DTypeLike
6
7__all__ = [
8    "atleast_1d",
9    "atleast_2d",
10    "atleast_3d",
11    "block",
12    "hstack",
13    "stack",
14    "unstack",
15    "vstack",
16]
17
18_T = TypeVar("_T")
19_ScalarT = TypeVar("_ScalarT", bound=generic)
20_ScalarT1 = TypeVar("_ScalarT1", bound=generic)
21_ScalarT2 = TypeVar("_ScalarT2", bound=generic)
22_ArrayT = TypeVar("_ArrayT", bound=NDArray[Any])
23
24###
25
26# keep in sync with `numpy.ma.extras.atleast_1d`
27@overload
28def atleast_1d(a0: _ArrayLike[_ScalarT], /) -> NDArray[_ScalarT]: ...
29@overload
30def atleast_1d(a0: _ArrayLike[_ScalarT1], a1: _ArrayLike[_ScalarT2], /) -> tuple[NDArray[_ScalarT1], NDArray[_ScalarT2]]: ...
31@overload
32def atleast_1d(a0: _ArrayLike[_ScalarT], a1: _ArrayLike[_ScalarT], /, *arys: _ArrayLike[_ScalarT]) -> tuple[NDArray[_ScalarT], ...]: ...
33@overload
34def atleast_1d(a0: ArrayLike, /) -> NDArray[Any]: ...
35@overload
36def atleast_1d(a0: ArrayLike, a1: ArrayLike, /) -> tuple[NDArray[Any], NDArray[Any]]: ...
37@overload
38def atleast_1d(a0: ArrayLike, a1: ArrayLike, /, *ai: ArrayLike) -> tuple[NDArray[Any], ...]: ...
39
40# keep in sync with `numpy.ma.extras.atleast_2d`
41@overload
42def atleast_2d(a0: _ArrayLike[_ScalarT], /) -> NDArray[_ScalarT]: ...
43@overload
44def atleast_2d(a0: _ArrayLike[_ScalarT1], a1: _ArrayLike[_ScalarT2], /) -> tuple[NDArray[_ScalarT1], NDArray[_ScalarT2]]: ...
45@overload
46def atleast_2d(a0: _ArrayLike[_ScalarT], a1: _ArrayLike[_ScalarT], /, *arys: _ArrayLike[_ScalarT]) -> tuple[NDArray[_ScalarT], ...]: ...
47@overload
48def atleast_2d(a0: ArrayLike, /) -> NDArray[Any]: ...
49@overload
50def atleast_2d(a0: ArrayLike, a1: ArrayLike, /) -> tuple[NDArray[Any], NDArray[Any]]: ...
51@overload
52def atleast_2d(a0: ArrayLike, a1: ArrayLike, /, *ai: ArrayLike) -> tuple[NDArray[Any], ...]: ...
53
54# keep in sync with `numpy.ma.extras.atleast_3d`
55@overload
56def atleast_3d(a0: _ArrayLike[_ScalarT], /) -> NDArray[_ScalarT]: ...
57@overload
58def atleast_3d(a0: _ArrayLike[_ScalarT1], a1: _ArrayLike[_ScalarT2], /) -> tuple[NDArray[_ScalarT1], NDArray[_ScalarT2]]: ...
59@overload
60def atleast_3d(a0: _ArrayLike[_ScalarT], a1: _ArrayLike[_ScalarT], /, *arys: _ArrayLike[_ScalarT]) -> tuple[NDArray[_ScalarT], ...]: ...
61@overload
62def atleast_3d(a0: ArrayLike, /) -> NDArray[Any]: ...
63@overload
64def atleast_3d(a0: ArrayLike, a1: ArrayLike, /) -> tuple[NDArray[Any], NDArray[Any]]: ...
65@overload
66def atleast_3d(a0: ArrayLike, a1: ArrayLike, /, *ai: ArrayLike) -> tuple[NDArray[Any], ...]: ...
67
68# used by numpy.lib._shape_base_impl
69def _arrays_for_stack_dispatcher(arrays: Sequence[_T]) -> tuple[_T, ...]: ...
70
71# keep in sync with `numpy.ma.extras.vstack`
72@overload
73def vstack(
74    tup: Sequence[_ArrayLike[_ScalarT]],
75    *,
76    dtype: None = None,
77    casting: _CastingKind = "same_kind"
78) -> NDArray[_ScalarT]: ...
79@overload
80def vstack(
81    tup: Sequence[ArrayLike],
82    *,
83    dtype: _DTypeLike[_ScalarT],
84    casting: _CastingKind = "same_kind"
85) -> NDArray[_ScalarT]: ...
86@overload
87def vstack(
88    tup: Sequence[ArrayLike],
89    *,
90    dtype: DTypeLike | None = None,
91    casting: _CastingKind = "same_kind"
92) -> NDArray[Any]: ...
93
94# keep in sync with `numpy.ma.extras.hstack`
95@overload
96def hstack(
97    tup: Sequence[_ArrayLike[_ScalarT]],
98    *,
99    dtype: None = None,
100    casting: _CastingKind = "same_kind"
101) -> NDArray[_ScalarT]: ...
102@overload
103def hstack(
104    tup: Sequence[ArrayLike],
105    *,
106    dtype: _DTypeLike[_ScalarT],
107    casting: _CastingKind = "same_kind"
108) -> NDArray[_ScalarT]: ...
109@overload
110def hstack(
111    tup: Sequence[ArrayLike],
112    *,
113    dtype: DTypeLike | None = None,
114    casting: _CastingKind = "same_kind"
115) -> NDArray[Any]: ...
116
117# keep in sync with `numpy.ma.extras.stack`
118@overload
119def stack(
120    arrays: Sequence[_ArrayLike[_ScalarT]],
121    axis: SupportsIndex = 0,
122    out: None = None,
123    *,
124    dtype: None = None,
125    casting: _CastingKind = "same_kind"
126) -> NDArray[_ScalarT]: ...
127@overload
128def stack(
129    arrays: Sequence[ArrayLike],
130    axis: SupportsIndex = 0,
131    out: None = None,
132    *,
133    dtype: _DTypeLike[_ScalarT],
134    casting: _CastingKind = "same_kind"
135) -> NDArray[_ScalarT]: ...
136@overload
137def stack(
138    arrays: Sequence[ArrayLike],
139    axis: SupportsIndex = 0,
140    out: None = None,
141    *,
142    dtype: DTypeLike | None = None,
143    casting: _CastingKind = "same_kind"
144) -> NDArray[Any]: ...
145@overload
146def stack(
147    arrays: Sequence[ArrayLike],
148    axis: SupportsIndex,
149    out: _ArrayT,
150    *,
151    dtype: DTypeLike | None = None,
152    casting: _CastingKind = "same_kind",
153) -> _ArrayT: ...
154@overload
155def stack(
156    arrays: Sequence[ArrayLike],
157    axis: SupportsIndex = 0,
158    *,
159    out: _ArrayT,
160    dtype: DTypeLike | None = None,
161    casting: _CastingKind = "same_kind",
162) -> _ArrayT: ...
163
164@overload
165def unstack(
166    array: _ArrayLike[_ScalarT],
167    /,
168    *,
169    axis: int = 0,
170) -> tuple[NDArray[_ScalarT], ...]: ...
171@overload
172def unstack(
173    array: ArrayLike,
174    /,
175    *,
176    axis: int = 0,
177) -> tuple[NDArray[Any], ...]: ...
178
179@overload
180def block(arrays: _ArrayLike[_ScalarT]) -> NDArray[_ScalarT]: ...
181@overload
182def block(arrays: ArrayLike) -> NDArray[Any]: ...
183