codekingpro/portable-devtools

1from collections.abc import Callable, Sequence
2from typing import (
3    Any,
4    Concatenate,
5    ParamSpec,
6    Protocol,
7    SupportsIndex,
8    TypeVar,
9    overload,
10    type_check_only,
11)
12from typing_extensions import deprecated
13
14import numpy as np
15from numpy import (
16    _CastingKind,
17    complexfloating,
18    floating,
19    generic,
20    integer,
21    object_,
22    signedinteger,
23    ufunc,
24    unsignedinteger,
25)
26from numpy._typing import (
27    ArrayLike,
28    DTypeLike,
29    NDArray,
30    _ArrayLike,
31    _ArrayLikeBool_co,
32    _ArrayLikeComplex_co,
33    _ArrayLikeFloat_co,
34    _ArrayLikeInt_co,
35    _ArrayLikeObject_co,
36    _ArrayLikeUInt_co,
37    _ShapeLike,
38)
39
40__all__ = [
41    "column_stack",
42    "row_stack",
43    "dstack",
44    "array_split",
45    "split",
46    "hsplit",
47    "vsplit",
48    "dsplit",
49    "apply_over_axes",
50    "expand_dims",
51    "apply_along_axis",
52    "kron",
53    "tile",
54    "take_along_axis",
55    "put_along_axis",
56]
57
58_P = ParamSpec("_P")
59_ScalarT = TypeVar("_ScalarT", bound=generic)
60
61# Signature of `__array_wrap__`
62@type_check_only
63class _ArrayWrap(Protocol):
64    def __call__(
65        self,
66        array: NDArray[Any],
67        context: tuple[ufunc, tuple[Any, ...], int] | None = ...,
68        return_scalar: bool = ...,
69        /,
70    ) -> Any: ...
71
72@type_check_only
73class _SupportsArrayWrap(Protocol):
74    @property
75    def __array_wrap__(self) -> _ArrayWrap: ...
76
77###
78
79def take_along_axis(
80    arr: _ScalarT | NDArray[_ScalarT],
81    indices: NDArray[integer],
82    axis: int | None = -1,
83) -> NDArray[_ScalarT]: ...
84
85def put_along_axis(
86    arr: NDArray[_ScalarT],
87    indices: NDArray[integer],
88    values: ArrayLike,
89    axis: int | None,
90) -> None: ...
91
92@overload
93def apply_along_axis(
94    func1d: Callable[Concatenate[NDArray[Any], _P], _ArrayLike[_ScalarT]],
95    axis: SupportsIndex,
96    arr: ArrayLike,
97    *args: _P.args,
98    **kwargs: _P.kwargs,
99) -> NDArray[_ScalarT]: ...
100@overload
101def apply_along_axis(
102    func1d: Callable[Concatenate[NDArray[Any], _P], Any],
103    axis: SupportsIndex,
104    arr: ArrayLike,
105    *args: _P.args,
106    **kwargs: _P.kwargs,
107) -> NDArray[Any]: ...
108
109def apply_over_axes(
110    func: Callable[[NDArray[Any], int], NDArray[_ScalarT]],
111    a: ArrayLike,
112    axes: int | Sequence[int],
113) -> NDArray[_ScalarT]: ...
114
115@overload
116def expand_dims(
117    a: _ArrayLike[_ScalarT],
118    axis: _ShapeLike,
119) -> NDArray[_ScalarT]: ...
120@overload
121def expand_dims(
122    a: ArrayLike,
123    axis: _ShapeLike,
124) -> NDArray[Any]: ...
125
126# Deprecated in NumPy 2.0, 2023-08-18
127@deprecated("`row_stack` alias is deprecated. Use `np.vstack` directly.")
128def row_stack(
129    tup: Sequence[ArrayLike],
130    *,
131    dtype: DTypeLike | None = None,
132    casting: _CastingKind = "same_kind",
133) -> NDArray[Any]: ...
134
135# keep in sync with `numpy.ma.extras.column_stack`
136@overload
137def column_stack(tup: Sequence[_ArrayLike[_ScalarT]]) -> NDArray[_ScalarT]: ...
138@overload
139def column_stack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ...
140
141# keep in sync with `numpy.ma.extras.dstack`
142@overload
143def dstack(tup: Sequence[_ArrayLike[_ScalarT]]) -> NDArray[_ScalarT]: ...
144@overload
145def dstack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ...
146
147@overload
148def array_split(
149    ary: _ArrayLike[_ScalarT],
150    indices_or_sections: _ShapeLike,
151    axis: SupportsIndex = 0,
152) -> list[NDArray[_ScalarT]]: ...
153@overload
154def array_split(
155    ary: ArrayLike,
156    indices_or_sections: _ShapeLike,
157    axis: SupportsIndex = 0,
158) -> list[NDArray[Any]]: ...
159
160@overload
161def split(
162    ary: _ArrayLike[_ScalarT],
163    indices_or_sections: _ShapeLike,
164    axis: SupportsIndex = 0,
165) -> list[NDArray[_ScalarT]]: ...
166@overload
167def split(
168    ary: ArrayLike,
169    indices_or_sections: _ShapeLike,
170    axis: SupportsIndex = 0,
171) -> list[NDArray[Any]]: ...
172
173# keep in sync with `numpy.ma.extras.hsplit`
174@overload
175def hsplit(
176    ary: _ArrayLike[_ScalarT],
177    indices_or_sections: _ShapeLike,
178) -> list[NDArray[_ScalarT]]: ...
179@overload
180def hsplit(
181    ary: ArrayLike,
182    indices_or_sections: _ShapeLike,
183) -> list[NDArray[Any]]: ...
184
185@overload
186def vsplit(
187    ary: _ArrayLike[_ScalarT],
188    indices_or_sections: _ShapeLike,
189) -> list[NDArray[_ScalarT]]: ...
190@overload
191def vsplit(
192    ary: ArrayLike,
193    indices_or_sections: _ShapeLike,
194) -> list[NDArray[Any]]: ...
195
196@overload
197def dsplit(
198    ary: _ArrayLike[_ScalarT],
199    indices_or_sections: _ShapeLike,
200) -> list[NDArray[_ScalarT]]: ...
201@overload
202def dsplit(
203    ary: ArrayLike,
204    indices_or_sections: _ShapeLike,
205) -> list[NDArray[Any]]: ...
206
207@overload
208def get_array_wrap(*args: _SupportsArrayWrap) -> _ArrayWrap: ...
209@overload
210def get_array_wrap(*args: object) -> _ArrayWrap | None: ...
211
212@overload
213def kron(a: _ArrayLikeBool_co, b: _ArrayLikeBool_co) -> NDArray[np.bool]: ...  # type: ignore[misc]
214@overload
215def kron(a: _ArrayLikeUInt_co, b: _ArrayLikeUInt_co) -> NDArray[unsignedinteger]: ...  # type: ignore[misc]
216@overload
217def kron(a: _ArrayLikeInt_co, b: _ArrayLikeInt_co) -> NDArray[signedinteger]: ...  # type: ignore[misc]
218@overload
219def kron(a: _ArrayLikeFloat_co, b: _ArrayLikeFloat_co) -> NDArray[floating]: ...  # type: ignore[misc]
220@overload
221def kron(a: _ArrayLikeComplex_co, b: _ArrayLikeComplex_co) -> NDArray[complexfloating]: ...
222@overload
223def kron(a: _ArrayLikeObject_co, b: Any) -> NDArray[object_]: ...
224@overload
225def kron(a: Any, b: _ArrayLikeObject_co) -> NDArray[object_]: ...
226
227@overload
228def tile(
229    A: _ArrayLike[_ScalarT],
230    reps: int | Sequence[int],
231) -> NDArray[_ScalarT]: ...
232@overload
233def tile(
234    A: ArrayLike,
235    reps: int | Sequence[int],
236) -> NDArray[Any]: ...
237