Skip to content

Commit 1057f6f

Browse files
jorenhamjorenham
authored
Merge pull request #768 from numpy:np24/ma.extras
👽️ `ma.extras` sync with upstream
2 parents 2e96342 + 8698ded commit 1057f6f

File tree

2 files changed

+185
-66
lines changed

2 files changed

+185
-66
lines changed

‎src/numpy-stubs/ma/extras.pyi‎

Lines changed: 185 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
from _typeshed import Incomplete
2-
from typing import Final, final
3-
from typing_extensions import override
2+
from collections.abc import Sequence
3+
from typing import Any, Final, SupportsIndex, final, overload
4+
from typing_extensions import TypeVar, override
45

56
import _numtype as _nt
7+
import numpy as np
8+
from numpy._typing import ArrayLike, DTypeLike, _ArrayLike, _DTypeLike, _ShapeLike
69
from numpy.lib._function_base_impl import average
710
from numpy.lib._index_tricks_impl import AxisConcatenator
811

@@ -57,27 +60,25 @@ __all__ = [
5760
"vstack",
5861
]
5962

60-
class _fromnxfunction:
61-
__name__: str
62-
__doc__: str
63-
def __init__(self, funcname: Incomplete) -> None: ...
64-
def getdoc(self) -> Incomplete: ...
65-
def __call__(self, *args: Incomplete, **params: Incomplete) -> Incomplete: ...
63+
_ScalarT = TypeVar("_ScalarT", bound=np.generic)
64+
_SCT0 = TypeVar("_SCT0", bound=np.generic)
65+
_SCT1 = TypeVar("_SCT1", bound=np.generic)
6666

67-
class _fromnxfunction_single(_fromnxfunction):
68-
__doc__: str
69-
@override
70-
def __call__(self, x: Incomplete, *args: Incomplete, **params: Incomplete) -> Incomplete: ...
67+
_ArrayT = TypeVar("_ArrayT", bound=_nt.Array[Any])
7168

72-
class _fromnxfunction_seq(_fromnxfunction):
73-
__doc__: str
74-
@override
75-
def __call__(self, x: Incomplete, *args: Incomplete, **params: Incomplete) -> Incomplete: ...
69+
_Array1T = TypeVar("_Array1T", bound=_nt.MArray[Any, _nt.Shape1N])
70+
_Array1T0 = TypeVar("_Array1T0", bound=_nt.MArray[Any, _nt.Shape1N])
71+
_Array1T1 = TypeVar("_Array1T1", bound=_nt.MArray[Any, _nt.Shape1N])
7672

77-
class _fromnxfunction_allargs(_fromnxfunction):
78-
__doc__: str
79-
@override
80-
def __call__(self, *args: Incomplete, **params: Incomplete) -> Incomplete: ...
73+
_Array2T = TypeVar("_Array2T", bound=_nt.MArray[Any, _nt.Shape2N])
74+
_Array2T0 = TypeVar("_Array2T0", bound=_nt.MArray[Any, _nt.Shape2N])
75+
_Array2T1 = TypeVar("_Array2T1", bound=_nt.MArray[Any, _nt.Shape2N])
76+
77+
_Array3T = TypeVar("_Array3T", bound=_nt.MArray[Any, _nt.Shape3N])
78+
_Array3T0 = TypeVar("_Array3T0", bound=_nt.MArray[Any, _nt.Shape3N])
79+
_Array3T1 = TypeVar("_Array3T1", bound=_nt.MArray[Any, _nt.Shape3N])
80+
81+
###
8182

8283
class MAxisConcatenator(AxisConcatenator):
8384
__slots__ = ()
@@ -133,14 +134,7 @@ def cov(
133134
allow_masked: Incomplete = ...,
134135
ddof: Incomplete = ...,
135136
) -> Incomplete: ...
136-
def corrcoef(
137-
x: Incomplete,
138-
y: Incomplete = ...,
139-
rowvar: Incomplete = ...,
140-
bias: Incomplete = ...,
141-
allow_masked: Incomplete = ...,
142-
ddof: Incomplete = ...,
143-
) -> Incomplete: ...
137+
def corrcoef(x: Incomplete, y: Incomplete = None, rowvar: bool = True, allow_masked: bool = True) -> Incomplete: ...
144138
def ndenumerate(a: Incomplete, compressed: Incomplete = ...) -> Incomplete: ...
145139
def flatnotmasked_edges(a: Incomplete) -> Incomplete: ...
146140
def notmasked_edges(a: Incomplete, axis: Incomplete = ...) -> Incomplete: ...
@@ -161,18 +155,170 @@ def polyfit(
161155

162156
mr_: Final[mr_class] = ...
163157

164-
atleast_1d: _fromnxfunction_allargs
165-
atleast_2d: _fromnxfunction_allargs
166-
atleast_3d: _fromnxfunction_allargs
158+
# keep in sync with `numpy._core.shape_base.atleast_1d`
159+
@overload
160+
def atleast_1d(a0: _Array1T, /) -> _Array1T: ...
161+
@overload
162+
def atleast_1d(a0: _Array1T0, a1: _Array1T1, /) -> tuple[_Array1T0, _Array1T1]: ...
163+
@overload
164+
def atleast_1d(a0: _Array1T, a1: _Array1T, /, *arys: _Array1T) -> tuple[_Array1T, ...]: ... # type: ignore[overload-overlap]
165+
@overload
166+
def atleast_1d(a0: _ArrayLike[_ScalarT], /) -> _nt.MArray[_ScalarT]: ...
167+
@overload
168+
def atleast_1d(a0: _ArrayLike[_SCT0], a2: _ArrayLike[_SCT1], /) -> tuple[_nt.MArray[_SCT0], _nt.MArray[_SCT1]]: ...
169+
@overload
170+
def atleast_1d(
171+
a0: _ArrayLike[_ScalarT], a2: _ArrayLike[_ScalarT], /, *arys: _ArrayLike[_ScalarT]
172+
) -> tuple[_nt.MArray[_ScalarT], ...]: ...
173+
@overload
174+
def atleast_1d(a0: ArrayLike, /) -> _nt.MArray[Any]: ...
175+
@overload
176+
def atleast_1d(a0: ArrayLike, a2: ArrayLike, /) -> tuple[_nt.MArray[Any], _nt.MArray[Any]]: ...
177+
@overload
178+
def atleast_1d(a0: ArrayLike, a2: ArrayLike, /, *arys: ArrayLike) -> tuple[_nt.MArray[Any], ...]: ...
179+
180+
# keep in sync with `numpy._core.shape_base.atleast_2d`
181+
@overload
182+
def atleast_2d(a0: _Array2T, /) -> _Array2T: ...
183+
@overload
184+
def atleast_2d(a0: _Array2T0, a1: _Array2T1, /) -> tuple[_Array2T0, _Array2T1]: ...
185+
@overload
186+
def atleast_2d(a0: _Array2T, a1: _Array2T, /, *arys: _Array2T) -> tuple[_Array2T, ...]: ... # type: ignore[overload-overlap]
187+
@overload
188+
def atleast_2d(a0: _ArrayLike[_ScalarT], /) -> _nt.MArray[_ScalarT]: ...
189+
@overload
190+
def atleast_2d(a0: _ArrayLike[_SCT0], a2: _ArrayLike[_SCT1], /) -> tuple[_nt.MArray[_SCT0], _nt.MArray[_SCT1]]: ...
191+
@overload
192+
def atleast_2d(
193+
a0: _ArrayLike[_ScalarT], a2: _ArrayLike[_ScalarT], /, *arys: _ArrayLike[_ScalarT]
194+
) -> tuple[_nt.MArray[_ScalarT], ...]: ...
195+
@overload
196+
def atleast_2d(a0: ArrayLike, /) -> _nt.MArray[Any]: ...
197+
@overload
198+
def atleast_2d(a0: ArrayLike, a2: ArrayLike, /) -> tuple[_nt.MArray[Any], _nt.MArray[Any]]: ...
199+
@overload
200+
def atleast_2d(a0: ArrayLike, a2: ArrayLike, /, *arys: ArrayLike) -> tuple[_nt.MArray[Any], ...]: ...
201+
202+
# keep in sync with `numpy._core.shape_base.atleast_2d`
203+
@overload
204+
def atleast_3d(a0: _Array3T, /) -> _Array3T: ...
205+
@overload
206+
def atleast_3d(a0: _Array3T0, a1: _Array3T1, /) -> tuple[_Array3T0, _Array3T1]: ...
207+
@overload
208+
def atleast_3d(a0: _Array3T, a1: _Array3T, /, *arys: _Array3T) -> tuple[_Array3T, ...]: ... # type: ignore[overload-overlap]
209+
@overload
210+
def atleast_3d(a0: _ArrayLike[_ScalarT], /) -> _nt.MArray[_ScalarT]: ...
211+
@overload
212+
def atleast_3d(a0: _ArrayLike[_SCT0], a2: _ArrayLike[_SCT1], /) -> tuple[_nt.MArray[_SCT0], _nt.MArray[_SCT1]]: ...
213+
@overload
214+
def atleast_3d(
215+
a0: _ArrayLike[_ScalarT], a2: _ArrayLike[_ScalarT], /, *arys: _ArrayLike[_ScalarT]
216+
) -> tuple[_nt.MArray[_ScalarT], ...]: ...
217+
@overload
218+
def atleast_3d(a0: ArrayLike, /) -> _nt.MArray[Any]: ...
219+
@overload
220+
def atleast_3d(a0: ArrayLike, a2: ArrayLike, /) -> tuple[_nt.MArray[Any], _nt.MArray[Any]]: ...
221+
@overload
222+
def atleast_3d(a0: ArrayLike, a2: ArrayLike, /, *arys: ArrayLike) -> tuple[_nt.MArray[Any], ...]: ...
223+
224+
# keep in sync with `numpy._core.shape_base.vstack`
225+
@overload
226+
def vstack(
227+
tup: Sequence[_ArrayLike[_ScalarT]], *, dtype: None = None, casting: np._CastingKind = "same_kind"
228+
) -> _nt.MArray[_ScalarT]: ...
229+
@overload
230+
def vstack(
231+
tup: Sequence[ArrayLike], *, dtype: _DTypeLike[_ScalarT], casting: np._CastingKind = "same_kind"
232+
) -> _nt.MArray[_ScalarT]: ...
233+
@overload
234+
def vstack(
235+
tup: Sequence[ArrayLike], *, dtype: DTypeLike | None = None, casting: np._CastingKind = "same_kind"
236+
) -> _nt.MArray[Any]: ...
237+
238+
row_stack = vstack
239+
240+
# keep in sync with `numpy._core.shape_base.hstack`
241+
@overload
242+
def hstack(
243+
tup: Sequence[_ArrayLike[_ScalarT]], *, dtype: None = None, casting: np._CastingKind = "same_kind"
244+
) -> _nt.MArray[_ScalarT]: ...
245+
@overload
246+
def hstack(
247+
tup: Sequence[ArrayLike], *, dtype: _DTypeLike[_ScalarT], casting: np._CastingKind = "same_kind"
248+
) -> _nt.MArray[_ScalarT]: ...
249+
@overload
250+
def hstack(
251+
tup: Sequence[ArrayLike], *, dtype: DTypeLike | None = None, casting: np._CastingKind = "same_kind"
252+
) -> _nt.MArray[Any]: ...
253+
254+
# keep in sync with `numpy._core.shape_base_impl.column_stack`
255+
@overload
256+
def column_stack(tup: Sequence[_ArrayLike[_ScalarT]]) -> _nt.MArray[_ScalarT]: ...
257+
@overload
258+
def column_stack(tup: Sequence[ArrayLike]) -> _nt.MArray[Incomplete]: ...
259+
260+
# keep in sync with `numpy._core.shape_base_impl.dstack`
261+
@overload
262+
def dstack(tup: Sequence[_ArrayLike[_ScalarT]]) -> _nt.MArray[_ScalarT]: ...
263+
@overload
264+
def dstack(tup: Sequence[ArrayLike]) -> _nt.MArray[Incomplete]: ...
167265

168-
vstack: _fromnxfunction_seq
169-
row_stack: _fromnxfunction_seq
170-
hstack: _fromnxfunction_seq
171-
column_stack: _fromnxfunction_seq
172-
dstack: _fromnxfunction_seq
173-
stack: _fromnxfunction_seq
266+
# keep in sync with `numpy._core.shape_base.stack`
267+
@overload
268+
def stack(
269+
arrays: Sequence[_ArrayLike[_ScalarT]],
270+
axis: SupportsIndex = 0,
271+
out: None = None,
272+
*,
273+
dtype: None = None,
274+
casting: np._CastingKind = "same_kind",
275+
) -> _nt.Array[_ScalarT]: ...
276+
@overload
277+
def stack(
278+
arrays: Sequence[ArrayLike],
279+
axis: SupportsIndex = 0,
280+
out: None = None,
281+
*,
282+
dtype: _DTypeLike[_ScalarT],
283+
casting: np._CastingKind = "same_kind",
284+
) -> _nt.Array[_ScalarT]: ...
285+
@overload
286+
def stack(
287+
arrays: Sequence[ArrayLike],
288+
axis: SupportsIndex = 0,
289+
out: None = None,
290+
*,
291+
dtype: DTypeLike | None = None,
292+
casting: np._CastingKind = "same_kind",
293+
) -> _nt.Array[Any]: ...
294+
@overload
295+
def stack(
296+
arrays: Sequence[ArrayLike],
297+
axis: SupportsIndex,
298+
out: _ArrayT,
299+
*,
300+
dtype: DTypeLike | None = None,
301+
casting: np._CastingKind = "same_kind",
302+
) -> _ArrayT: ...
303+
@overload
304+
def stack(
305+
arrays: Sequence[ArrayLike],
306+
axis: SupportsIndex = 0,
307+
*,
308+
out: _ArrayT,
309+
dtype: DTypeLike | None = None,
310+
casting: np._CastingKind = "same_kind",
311+
) -> _ArrayT: ...
174312

175-
hsplit: _fromnxfunction_single
176-
diagflat: _fromnxfunction_single
313+
# keep in sync with `numpy._core.shape_base_impl.hsplit`
314+
@overload
315+
def hsplit(ary: _ArrayLike[_ScalarT], indices_or_sections: _ShapeLike) -> list[_nt.MArray[_ScalarT]]: ...
316+
@overload
317+
def hsplit(ary: ArrayLike, indices_or_sections: _ShapeLike) -> list[_nt.MArray[Incomplete]]: ...
177318

319+
# keep in sync with `numpy._core.twodim_base_impl.hsplit`
320+
@overload
321+
def diagflat(v: _ArrayLike[_ScalarT], k: int = 0) -> _nt.MArray[_ScalarT]: ...
322+
@overload
323+
def diagflat(v: ArrayLike, k: int = 0) -> _nt.MArray[Incomplete]: ...
178324
def mask_rowcols(a: Incomplete, axis: Incomplete | None = None) -> _nt.MArray[Incomplete]: ...

‎tool/allowlists/todo.txt‎

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,7 @@ numpy.ma.arange
66
numpy.ma.argmax
77
numpy.ma.argmin
88
numpy.ma.asanyarray
9-
numpy.ma.atleast_1d
10-
numpy.ma.atleast_2d
11-
numpy.ma.atleast_3d
129
numpy.ma.clip
13-
numpy.ma.column_stack
1410
numpy.ma.compress
1511
numpy.ma.copy
1612
numpy.ma.core._convert2ma
@@ -55,35 +51,15 @@ numpy.ma.core.trace
5551
numpy.ma.core.var
5652
numpy.ma.core.zeros
5753
numpy.ma.core.zeros_like
58-
numpy.ma.corrcoef
5954
numpy.ma.count
6055
numpy.ma.cumprod
6156
numpy.ma.cumsum
62-
numpy.ma.diagflat
6357
numpy.ma.diagonal
64-
numpy.ma.dstack
6558
numpy.ma.empty
6659
numpy.ma.empty_like
67-
numpy.ma.extras._fromnxfunction_allargs
68-
numpy.ma.extras._fromnxfunction_seq
69-
numpy.ma.extras._fromnxfunction_single
70-
numpy.ma.extras.atleast_1d
71-
numpy.ma.extras.atleast_2d
72-
numpy.ma.extras.atleast_3d
73-
numpy.ma.extras.column_stack
74-
numpy.ma.extras.corrcoef
75-
numpy.ma.extras.diagflat
76-
numpy.ma.extras.dstack
77-
numpy.ma.extras.hsplit
78-
numpy.ma.extras.hstack
79-
numpy.ma.extras.row_stack
80-
numpy.ma.extras.stack
81-
numpy.ma.extras.vstack
8260
numpy.ma.frombuffer
8361
numpy.ma.fromfunction
8462
numpy.ma.harden_mask
85-
numpy.ma.hsplit
86-
numpy.ma.hstack
8763
numpy.ma.identity
8864
numpy.ma.ids
8965
numpy.ma.indices
@@ -96,15 +72,12 @@ numpy.ma.prod
9672
numpy.ma.product
9773
numpy.ma.ravel
9874
numpy.ma.repeat
99-
numpy.ma.row_stack
10075
numpy.ma.soften_mask
10176
numpy.ma.squeeze
102-
numpy.ma.stack
10377
numpy.ma.std
10478
numpy.ma.sum
10579
numpy.ma.swapaxes
10680
numpy.ma.trace
10781
numpy.ma.var
108-
numpy.ma.vstack
10982
numpy.ma.zeros
11083
numpy.ma.zeros_like

0 commit comments

Comments
 (0)