codekingpro/portable-devtools

1from _typeshed import Incomplete, SupportsLenAndGetItem
2from collections.abc import Sequence
3from typing import (
4    Any,
5    ClassVar,
6    Final,
7    Generic,
8    Literal as L,
9    Self,
10    SupportsIndex,
11    final,
12    overload,
13)
14from typing_extensions import TypeVar
15
16import numpy as np
17from numpy import _CastingKind
18from numpy._core.multiarray import ravel_multi_index, unravel_index
19from numpy._typing import (
20    ArrayLike,
21    DTypeLike,
22    NDArray,
23    _AnyShape,
24    _ArrayLike,
25    _DTypeLike,
26    _FiniteNestedSequence,
27    _HasDType,
28    _NestedSequence,
29    _SupportsArray,
30)
31
32__all__ = [  # noqa: RUF022
33    "ravel_multi_index",
34    "unravel_index",
35    "mgrid",
36    "ogrid",
37    "r_",
38    "c_",
39    "s_",
40    "index_exp",
41    "ix_",
42    "ndenumerate",
43    "ndindex",
44    "fill_diagonal",
45    "diag_indices",
46    "diag_indices_from",
47]
48
49###
50
51_T = TypeVar("_T")
52_TupleT = TypeVar("_TupleT", bound=tuple[Any, ...])
53_ArrayT = TypeVar("_ArrayT", bound=NDArray[Any])
54_DTypeT = TypeVar("_DTypeT", bound=np.dtype)
55_ScalarT = TypeVar("_ScalarT", bound=np.generic)
56_ScalarT_co = TypeVar("_ScalarT_co", bound=np.generic, default=Any, covariant=True)
57_BoolT_co = TypeVar("_BoolT_co", bound=bool, default=bool, covariant=True)
58
59_AxisT_co = TypeVar("_AxisT_co", bound=int, default=L[0], covariant=True)
60_MatrixT_co = TypeVar("_MatrixT_co", bound=bool, default=L[False], covariant=True)
61_NDMinT_co = TypeVar("_NDMinT_co", bound=int, default=L[1], covariant=True)
62_Trans1DT_co = TypeVar("_Trans1DT_co", bound=int, default=L[-1], covariant=True)
63
64###
65
66class ndenumerate(Generic[_ScalarT_co]):
67    @overload
68    def __init__(self: ndenumerate[_ScalarT], arr: _FiniteNestedSequence[_SupportsArray[np.dtype[_ScalarT]]]) -> None: ...
69    @overload
70    def __init__(self: ndenumerate[np.str_], arr: str | _NestedSequence[str]) -> None: ...
71    @overload
72    def __init__(self: ndenumerate[np.bytes_], arr: bytes | _NestedSequence[bytes]) -> None: ...
73    @overload
74    def __init__(self: ndenumerate[np.bool], arr: bool | _NestedSequence[bool]) -> None: ...
75    @overload
76    def __init__(self: ndenumerate[np.intp], arr: int | _NestedSequence[int]) -> None: ...
77    @overload
78    def __init__(self: ndenumerate[np.float64], arr: float | _NestedSequence[float]) -> None: ...
79    @overload
80    def __init__(self: ndenumerate[np.complex128], arr: complex | _NestedSequence[complex]) -> None: ...
81    @overload
82    def __init__(self: ndenumerate[Incomplete], arr: object) -> None: ...
83
84    # The first overload is a (semi-)workaround for a mypy bug (tested with v1.10 and v1.11)
85    @overload
86    def __next__(
87        self: ndenumerate[np.bool | np.number | np.flexible | np.datetime64 | np.timedelta64],
88        /,
89    ) -> tuple[_AnyShape, _ScalarT_co]: ...
90    @overload
91    def __next__(self: ndenumerate[np.object_], /) -> tuple[_AnyShape, Incomplete]: ...
92    @overload
93    def __next__(self, /) -> tuple[_AnyShape, _ScalarT_co]: ...
94
95    #
96    def __iter__(self) -> Self: ...
97
98class ndindex:
99    @overload
100    def __init__(self, shape: tuple[SupportsIndex, ...], /) -> None: ...
101    @overload
102    def __init__(self, /, *shape: SupportsIndex) -> None: ...
103
104    #
105    def __iter__(self) -> Self: ...
106    def __next__(self) -> _AnyShape: ...
107
108class nd_grid(Generic[_BoolT_co]):
109    __slots__ = ("sparse",)
110
111    sparse: _BoolT_co
112    def __init__(self, sparse: _BoolT_co = ...) -> None: ...  # stubdefaulter: ignore[missing-default]
113    @overload
114    def __getitem__(self: nd_grid[L[False]], key: slice | Sequence[slice]) -> NDArray[Incomplete]: ...
115    @overload
116    def __getitem__(self: nd_grid[L[True]], key: slice | Sequence[slice]) -> tuple[NDArray[Incomplete], ...]: ...
117
118@final
119class MGridClass(nd_grid[L[False]]):
120    __slots__ = ()
121
122    def __init__(self) -> None: ...
123
124@final
125class OGridClass(nd_grid[L[True]]):
126    __slots__ = ()
127
128    def __init__(self) -> None: ...
129
130class AxisConcatenator(Generic[_AxisT_co, _MatrixT_co, _NDMinT_co, _Trans1DT_co]):
131    __slots__ = "axis", "matrix", "ndmin", "trans1d"
132
133    makemat: ClassVar[type[np.matrix[tuple[int, int], np.dtype]]]
134
135    axis: _AxisT_co
136    matrix: _MatrixT_co
137    ndmin: _NDMinT_co
138    trans1d: _Trans1DT_co
139
140    # NOTE: mypy does not understand that these default values are the same as the
141    # TypeVar defaults. Since the workaround would require us to write 16 overloads,
142    # we ignore the assignment type errors here.
143    def __init__(
144        self,
145        /,
146        axis: _AxisT_co = 0,  # type: ignore[assignment]
147        matrix: _MatrixT_co = False,  # type: ignore[assignment]
148        ndmin: _NDMinT_co = 1,  # type: ignore[assignment]
149        trans1d: _Trans1DT_co = -1,  # type: ignore[assignment]
150    ) -> None: ...
151
152    # TODO(jorenham): annotate this
153    def __getitem__(self, key: Incomplete, /) -> Incomplete: ...
154    def __len__(self, /) -> L[0]: ...
155
156    # Keep in sync with _core.multiarray.concatenate
157    @staticmethod
158    @overload
159    def concatenate(
160        arrays: _ArrayLike[_ScalarT],
161        /,
162        axis: SupportsIndex | None = 0,
163        out: None = None,
164        *,
165        dtype: None = None,
166        casting: _CastingKind | None = "same_kind",
167    ) -> NDArray[_ScalarT]: ...
168    @staticmethod
169    @overload
170    def concatenate(
171        arrays: SupportsLenAndGetItem[ArrayLike],
172        /,
173        axis: SupportsIndex | None = 0,
174        out: None = None,
175        *,
176        dtype: _DTypeLike[_ScalarT],
177        casting: _CastingKind | None = "same_kind",
178    ) -> NDArray[_ScalarT]: ...
179    @staticmethod
180    @overload
181    def concatenate(
182        arrays: SupportsLenAndGetItem[ArrayLike],
183        /,
184        axis: SupportsIndex | None = 0,
185        out: None = None,
186        *,
187        dtype: DTypeLike | None = None,
188        casting: _CastingKind | None = "same_kind",
189    ) -> NDArray[Incomplete]: ...
190    @staticmethod
191    @overload
192    def concatenate(
193        arrays: SupportsLenAndGetItem[ArrayLike],
194        /,
195        axis: SupportsIndex | None = 0,
196        *,
197        out: _ArrayT,
198        dtype: DTypeLike | None = None,
199        casting: _CastingKind | None = "same_kind",
200    ) -> _ArrayT: ...
201    @staticmethod
202    @overload
203    def concatenate(
204        arrays: SupportsLenAndGetItem[ArrayLike],
205        /,
206        axis: SupportsIndex | None,
207        out: _ArrayT,
208        *,
209        dtype: DTypeLike | None = None,
210        casting: _CastingKind | None = "same_kind",
211    ) -> _ArrayT: ...
212
213@final
214class RClass(AxisConcatenator[L[0], L[False], L[1], L[-1]]):
215    __slots__ = ()
216
217    def __init__(self, /) -> None: ...
218
219@final
220class CClass(AxisConcatenator[L[-1], L[False], L[2], L[0]]):
221    __slots__ = ()
222
223    def __init__(self, /) -> None: ...
224
225class IndexExpression(Generic[_BoolT_co]):
226    __slots__ = ("maketuple",)
227
228    maketuple: _BoolT_co
229    def __init__(self, maketuple: _BoolT_co) -> None: ...
230    @overload
231    def __getitem__(self, item: _TupleT) -> _TupleT: ...
232    @overload
233    def __getitem__(self: IndexExpression[L[True]], item: _T) -> tuple[_T]: ...
234    @overload
235    def __getitem__(self: IndexExpression[L[False]], item: _T) -> _T: ...
236
237@overload
238def ix_(*args: _FiniteNestedSequence[_HasDType[_DTypeT]]) -> tuple[np.ndarray[_AnyShape, _DTypeT], ...]: ...
239@overload
240def ix_(*args: str | _NestedSequence[str]) -> tuple[NDArray[np.str_], ...]: ...
241@overload
242def ix_(*args: bytes | _NestedSequence[bytes]) -> tuple[NDArray[np.bytes_], ...]: ...
243@overload
244def ix_(*args: bool | _NestedSequence[bool]) -> tuple[NDArray[np.bool], ...]: ...
245@overload
246def ix_(*args: int | _NestedSequence[int]) -> tuple[NDArray[np.intp], ...]: ...
247@overload
248def ix_(*args: float | _NestedSequence[float]) -> tuple[NDArray[np.float64], ...]: ...
249@overload
250def ix_(*args: complex | _NestedSequence[complex]) -> tuple[NDArray[np.complex128], ...]: ...
251
252#
253def fill_diagonal(a: NDArray[Any], val: object, wrap: bool = False) -> None: ...
254
255#
256def diag_indices(n: int, ndim: int = 2) -> tuple[NDArray[np.intp], ...]: ...
257def diag_indices_from(arr: ArrayLike) -> tuple[NDArray[np.intp], ...]: ...
258
259#
260mgrid: Final[MGridClass] = ...
261ogrid: Final[OGridClass] = ...
262
263r_: Final[RClass] = ...
264c_: Final[CClass] = ...
265
266index_exp: Final[IndexExpression[L[True]]] = ...
267s_: Final[IndexExpression[L[False]]] = ...
268