codekingpro/portable-devtools

1from _typeshed import Incomplete
2from collections.abc import Sequence
3from typing import SupportsIndex, TypeAlias, TypeVar, overload
4
5import numpy as np
6from numpy import _CastingKind
7from numpy._typing import (
8    ArrayLike,
9    DTypeLike,
10    _AnyShape,
11    _ArrayLike,
12    _DTypeLike,
13    _ShapeLike,
14)
15from numpy.lib._function_base_impl import average
16from numpy.lib._index_tricks_impl import AxisConcatenator
17
18from .core import MaskedArray, dot
19
20__all__ = [
21    "apply_along_axis",
22    "apply_over_axes",
23    "atleast_1d",
24    "atleast_2d",
25    "atleast_3d",
26    "average",
27    "clump_masked",
28    "clump_unmasked",
29    "column_stack",
30    "compress_cols",
31    "compress_nd",
32    "compress_rowcols",
33    "compress_rows",
34    "corrcoef",
35    "count_masked",
36    "cov",
37    "diagflat",
38    "dot",
39    "dstack",
40    "ediff1d",
41    "flatnotmasked_contiguous",
42    "flatnotmasked_edges",
43    "hsplit",
44    "hstack",
45    "in1d",
46    "intersect1d",
47    "isin",
48    "mask_cols",
49    "mask_rowcols",
50    "mask_rows",
51    "masked_all",
52    "masked_all_like",
53    "median",
54    "mr_",
55    "ndenumerate",
56    "notmasked_contiguous",
57    "notmasked_edges",
58    "polyfit",
59    "row_stack",
60    "setdiff1d",
61    "setxor1d",
62    "stack",
63    "union1d",
64    "unique",
65    "vander",
66    "vstack",
67]
68
69_ScalarT = TypeVar("_ScalarT", bound=np.generic)
70_ScalarT1 = TypeVar("_ScalarT1", bound=np.generic)
71_ScalarT2 = TypeVar("_ScalarT2", bound=np.generic)
72_MArrayT = TypeVar("_MArrayT", bound=MaskedArray)
73
74_MArray: TypeAlias = MaskedArray[_AnyShape, np.dtype[_ScalarT]]
75
76###
77
78# keep in sync with `numpy._core.shape_base.atleast_1d`
79@overload
80def atleast_1d(a0: _ArrayLike[_ScalarT], /) -> _MArray[_ScalarT]: ...
81@overload
82def atleast_1d(a0: _ArrayLike[_ScalarT1], a1: _ArrayLike[_ScalarT2], /) -> tuple[_MArray[_ScalarT1], _MArray[_ScalarT2]]: ...
83@overload
84def atleast_1d(
85    a0: _ArrayLike[_ScalarT], a1: _ArrayLike[_ScalarT], /, *arys: _ArrayLike[_ScalarT]
86) -> tuple[_MArray[_ScalarT], ...]: ...
87@overload
88def atleast_1d(a0: ArrayLike, /) -> _MArray[Incomplete]: ...
89@overload
90def atleast_1d(a0: ArrayLike, a1: ArrayLike, /) -> tuple[_MArray[Incomplete], _MArray[Incomplete]]: ...
91@overload
92def atleast_1d(a0: ArrayLike, a1: ArrayLike, /, *ai: ArrayLike) -> tuple[_MArray[Incomplete], ...]: ...
93
94# keep in sync with `numpy._core.shape_base.atleast_2d`
95@overload
96def atleast_2d(a0: _ArrayLike[_ScalarT], /) -> _MArray[_ScalarT]: ...
97@overload
98def atleast_2d(a0: _ArrayLike[_ScalarT1], a1: _ArrayLike[_ScalarT2], /) -> tuple[_MArray[_ScalarT1], _MArray[_ScalarT2]]: ...
99@overload
100def atleast_2d(
101    a0: _ArrayLike[_ScalarT], a1: _ArrayLike[_ScalarT], /, *arys: _ArrayLike[_ScalarT]
102) -> tuple[_MArray[_ScalarT], ...]: ...
103@overload
104def atleast_2d(a0: ArrayLike, /) -> _MArray[Incomplete]: ...
105@overload
106def atleast_2d(a0: ArrayLike, a1: ArrayLike, /) -> tuple[_MArray[Incomplete], _MArray[Incomplete]]: ...
107@overload
108def atleast_2d(a0: ArrayLike, a1: ArrayLike, /, *ai: ArrayLike) -> tuple[_MArray[Incomplete], ...]: ...
109
110# keep in sync with `numpy._core.shape_base.atleast_2d`
111@overload
112def atleast_3d(a0: _ArrayLike[_ScalarT], /) -> _MArray[_ScalarT]: ...
113@overload
114def atleast_3d(a0: _ArrayLike[_ScalarT1], a1: _ArrayLike[_ScalarT2], /) -> tuple[_MArray[_ScalarT1], _MArray[_ScalarT2]]: ...
115@overload
116def atleast_3d(
117    a0: _ArrayLike[_ScalarT], a1: _ArrayLike[_ScalarT], /, *arys: _ArrayLike[_ScalarT]
118) -> tuple[_MArray[_ScalarT], ...]: ...
119@overload
120def atleast_3d(a0: ArrayLike, /) -> _MArray[Incomplete]: ...
121@overload
122def atleast_3d(a0: ArrayLike, a1: ArrayLike, /) -> tuple[_MArray[Incomplete], _MArray[Incomplete]]: ...
123@overload
124def atleast_3d(a0: ArrayLike, a1: ArrayLike, /, *ai: ArrayLike) -> tuple[_MArray[Incomplete], ...]: ...
125
126# keep in sync with `numpy._core.shape_base.vstack`
127@overload
128def vstack(
129    tup: Sequence[_ArrayLike[_ScalarT]],
130    *,
131    dtype: None = None,
132    casting: _CastingKind = "same_kind"
133) -> _MArray[_ScalarT]: ...
134@overload
135def vstack(
136    tup: Sequence[ArrayLike],
137    *,
138    dtype: _DTypeLike[_ScalarT],
139    casting: _CastingKind = "same_kind"
140) -> _MArray[_ScalarT]: ...
141@overload
142def vstack(
143    tup: Sequence[ArrayLike],
144    *,
145    dtype: DTypeLike | None = None,
146    casting: _CastingKind = "same_kind"
147) -> _MArray[Incomplete]: ...
148
149row_stack = vstack
150
151# keep in sync with `numpy._core.shape_base.hstack`
152@overload
153def hstack(
154    tup: Sequence[_ArrayLike[_ScalarT]],
155    *,
156    dtype: None = None,
157    casting: _CastingKind = "same_kind"
158) -> _MArray[_ScalarT]: ...
159@overload
160def hstack(
161    tup: Sequence[ArrayLike],
162    *,
163    dtype: _DTypeLike[_ScalarT],
164    casting: _CastingKind = "same_kind"
165) -> _MArray[_ScalarT]: ...
166@overload
167def hstack(
168    tup: Sequence[ArrayLike],
169    *,
170    dtype: DTypeLike | None = None,
171    casting: _CastingKind = "same_kind"
172) -> _MArray[Incomplete]: ...
173
174# keep in sync with `numpy._core.shape_base_impl.column_stack`
175@overload
176def column_stack(tup: Sequence[_ArrayLike[_ScalarT]]) -> _MArray[_ScalarT]: ...
177@overload
178def column_stack(tup: Sequence[ArrayLike]) -> _MArray[Incomplete]: ...
179
180# keep in sync with `numpy._core.shape_base_impl.dstack`
181@overload
182def dstack(tup: Sequence[_ArrayLike[_ScalarT]]) -> _MArray[_ScalarT]: ...
183@overload
184def dstack(tup: Sequence[ArrayLike]) -> _MArray[Incomplete]: ...
185
186# keep in sync with `numpy._core.shape_base.stack`
187@overload
188def stack(
189    arrays: Sequence[_ArrayLike[_ScalarT]],
190    axis: SupportsIndex = 0,
191    out: None = None,
192    *,
193    dtype: None = None,
194    casting: _CastingKind = "same_kind"
195) -> _MArray[_ScalarT]: ...
196@overload
197def stack(
198    arrays: Sequence[ArrayLike],
199    axis: SupportsIndex = 0,
200    out: None = None,
201    *,
202    dtype: _DTypeLike[_ScalarT],
203    casting: _CastingKind = "same_kind"
204) -> _MArray[_ScalarT]: ...
205@overload
206def stack(
207    arrays: Sequence[ArrayLike],
208    axis: SupportsIndex = 0,
209    out: None = None,
210    *,
211    dtype: DTypeLike | None = None,
212    casting: _CastingKind = "same_kind"
213) -> _MArray[Incomplete]: ...
214@overload
215def stack(
216    arrays: Sequence[ArrayLike],
217    axis: SupportsIndex,
218    out: _MArrayT,
219    *,
220    dtype: DTypeLike | None = None,
221    casting: _CastingKind = "same_kind",
222) -> _MArrayT: ...
223@overload
224def stack(
225    arrays: Sequence[ArrayLike],
226    axis: SupportsIndex = 0,
227    *,
228    out: _MArrayT,
229    dtype: DTypeLike | None = None,
230    casting: _CastingKind = "same_kind",
231) -> _MArrayT: ...
232
233# keep in sync with `numpy._core.shape_base_impl.hsplit`
234@overload
235def hsplit(ary: _ArrayLike[_ScalarT], indices_or_sections: _ShapeLike) -> list[_MArray[_ScalarT]]: ...
236@overload
237def hsplit(ary: ArrayLike, indices_or_sections: _ShapeLike) -> list[_MArray[Incomplete]]: ...
238
239# keep in sync with `numpy._core.twodim_base_impl.hsplit`
240@overload
241def diagflat(v: _ArrayLike[_ScalarT], k: int = 0) -> _MArray[_ScalarT]: ...
242@overload
243def diagflat(v: ArrayLike, k: int = 0) -> _MArray[Incomplete]: ...
244
245# TODO: everything below
246
247def count_masked(arr, axis=None): ...
248def masked_all(shape, dtype=float): ...  # noqa: PYI014
249def masked_all_like(arr): ...
250
251def apply_along_axis(func1d, axis, arr, *args, **kwargs): ...
252def apply_over_axes(func, a, axes): ...
253def median(a, axis=None, out=None, overwrite_input=False, keepdims=False): ...
254def compress_nd(x, axis=None): ...
255def compress_rowcols(x, axis=None): ...
256def compress_rows(a): ...
257def compress_cols(a): ...
258def mask_rows(a, axis=...): ...
259def mask_cols(a, axis=...): ...
260def ediff1d(arr, to_end=None, to_begin=None): ...
261def unique(ar1, return_index=False, return_inverse=False): ...
262def intersect1d(ar1, ar2, assume_unique=False): ...
263def setxor1d(ar1, ar2, assume_unique=False): ...
264def in1d(ar1, ar2, assume_unique=False, invert=False): ...
265def isin(element, test_elements, assume_unique=False, invert=False): ...
266def union1d(ar1, ar2): ...
267def setdiff1d(ar1, ar2, assume_unique=False): ...
268def cov(x, y=None, rowvar=True, bias=False, allow_masked=True, ddof=None): ...
269def corrcoef(x, y=None, rowvar=True, allow_masked=True): ...
270
271class MAxisConcatenator(AxisConcatenator):
272    __slots__ = ()
273
274    @staticmethod
275    def concatenate(arrays: Incomplete, axis: int = 0) -> Incomplete: ...  # type: ignore[override]  # pyright: ignore[reportIncompatibleMethodOverride]
276    @classmethod
277    def makemat(cls, arr: Incomplete) -> Incomplete: ...  # type: ignore[override]  # pyright: ignore[reportIncompatibleVariableOverride]
278
279class mr_class(MAxisConcatenator):
280    __slots__ = ()
281
282    def __init__(self) -> None: ...
283
284mr_: mr_class
285
286def ndenumerate(a, compressed=True): ...
287def flatnotmasked_edges(a): ...
288def notmasked_edges(a, axis=None): ...
289def flatnotmasked_contiguous(a): ...
290def notmasked_contiguous(a, axis=None): ...
291def clump_unmasked(a): ...
292def clump_masked(a): ...
293def vander(x, n=None): ...
294def polyfit(x, y, deg, rcond=None, full=False, w=None, cov=False): ...
295
296#
297def mask_rowcols(a: Incomplete, axis: Incomplete | None = None) -> MaskedArray[Incomplete, np.dtype[Incomplete]]: ...
298