👽️ ma.extras sync with upstream

src/numpy-stubs/ma/extras.pyi

@@ -1,8 +1,11 @@
 from _typeshed import Incomplete
-from typing import Final, final
-from typing_extensions import override
+from collections.abc import Sequence
+from typing import Any, Final, SupportsIndex, final, overload
+from typing_extensions import TypeVar, override
 
 import _numtype as _nt
+import numpy as np
+from numpy._typing import ArrayLike, DTypeLike, _ArrayLike, _DTypeLike, _ShapeLike
 from numpy.lib._function_base_impl import average
 from numpy.lib._index_tricks_impl import AxisConcatenator
 
@@ -57,27 +60,25 @@ __all__ = [
     "vstack",
 ]
 
-class _fromnxfunction:
-    __name__: str
-    __doc__: str
-    def __init__(self, funcname: Incomplete) -> None: ...
-    def getdoc(self) -> Incomplete: ...
-    def __call__(self, *args: Incomplete, **params: Incomplete) -> Incomplete: ...
+_ScalarT = TypeVar("_ScalarT", bound=np.generic)
+_SCT0 = TypeVar("_SCT0", bound=np.generic)
+_SCT1 = TypeVar("_SCT1", bound=np.generic)
 
-class _fromnxfunction_single(_fromnxfunction):
-    __doc__: str
-    @override
-    def __call__(self, x: Incomplete, *args: Incomplete, **params: Incomplete) -> Incomplete: ...
+_ArrayT = TypeVar("_ArrayT", bound=_nt.Array[Any])
 
-class _fromnxfunction_seq(_fromnxfunction):
-    __doc__: str
-    @override
-    def __call__(self, x: Incomplete, *args: Incomplete, **params: Incomplete) -> Incomplete: ...
+_Array1T = TypeVar("_Array1T", bound=_nt.MArray[Any, _nt.Shape1N])
+_Array1T0 = TypeVar("_Array1T0", bound=_nt.MArray[Any, _nt.Shape1N])
+_Array1T1 = TypeVar("_Array1T1", bound=_nt.MArray[Any, _nt.Shape1N])
 
-class _fromnxfunction_allargs(_fromnxfunction):
-    __doc__: str
-    @override
-    def __call__(self, *args: Incomplete, **params: Incomplete) -> Incomplete: ...
+_Array2T = TypeVar("_Array2T", bound=_nt.MArray[Any, _nt.Shape2N])
+_Array2T0 = TypeVar("_Array2T0", bound=_nt.MArray[Any, _nt.Shape2N])
+_Array2T1 = TypeVar("_Array2T1", bound=_nt.MArray[Any, _nt.Shape2N])
+
+_Array3T = TypeVar("_Array3T", bound=_nt.MArray[Any, _nt.Shape3N])
+_Array3T0 = TypeVar("_Array3T0", bound=_nt.MArray[Any, _nt.Shape3N])
+_Array3T1 = TypeVar("_Array3T1", bound=_nt.MArray[Any, _nt.Shape3N])
+
+###
 
 class MAxisConcatenator(AxisConcatenator):
     __slots__ = ()
@@ -133,14 +134,7 @@ def cov(
     allow_masked: Incomplete = ...,
     ddof: Incomplete = ...,
 ) -> Incomplete: ...
-def corrcoef(
-    x: Incomplete,
-    y: Incomplete = ...,
-    rowvar: Incomplete = ...,
-    bias: Incomplete = ...,
-    allow_masked: Incomplete = ...,
-    ddof: Incomplete = ...,
-) -> Incomplete: ...
+def corrcoef(x: Incomplete, y: Incomplete = None, rowvar: bool = True, allow_masked: bool = True) -> Incomplete: ...
 def ndenumerate(a: Incomplete, compressed: Incomplete = ...) -> Incomplete: ...
 def flatnotmasked_edges(a: Incomplete) -> Incomplete: ...
 def notmasked_edges(a: Incomplete, axis: Incomplete = ...) -> Incomplete: ...
@@ -161,18 +155,170 @@ def polyfit(
 
 mr_: Final[mr_class] = ...
 
-atleast_1d: _fromnxfunction_allargs
-atleast_2d: _fromnxfunction_allargs
-atleast_3d: _fromnxfunction_allargs
+# keep in sync with `numpy._core.shape_base.atleast_1d`
+@overload
+def atleast_1d(a0: _Array1T, /) -> _Array1T: ...
+@overload
+def atleast_1d(a0: _Array1T0, a1: _Array1T1, /) -> tuple[_Array1T0, _Array1T1]: ...
+@overload
+def atleast_1d(a0: _Array1T, a1: _Array1T, /, *arys: _Array1T) -> tuple[_Array1T, ...]: ...  # type: ignore[overload-overlap]
+@overload
+def atleast_1d(a0: _ArrayLike[_ScalarT], /) -> _nt.MArray[_ScalarT]: ...
+@overload
+def atleast_1d(a0: _ArrayLike[_SCT0], a2: _ArrayLike[_SCT1], /) -> tuple[_nt.MArray[_SCT0], _nt.MArray[_SCT1]]: ...
+@overload
+def atleast_1d(
+    a0: _ArrayLike[_ScalarT], a2: _ArrayLike[_ScalarT], /, *arys: _ArrayLike[_ScalarT]
+) -> tuple[_nt.MArray[_ScalarT], ...]: ...
+@overload
+def atleast_1d(a0: ArrayLike, /) -> _nt.MArray[Any]: ...
+@overload
+def atleast_1d(a0: ArrayLike, a2: ArrayLike, /) -> tuple[_nt.MArray[Any], _nt.MArray[Any]]: ...
+@overload
+def atleast_1d(a0: ArrayLike, a2: ArrayLike, /, *arys: ArrayLike) -> tuple[_nt.MArray[Any], ...]: ...
+
+# keep in sync with `numpy._core.shape_base.atleast_2d`
+@overload
+def atleast_2d(a0: _Array2T, /) -> _Array2T: ...
+@overload
+def atleast_2d(a0: _Array2T0, a1: _Array2T1, /) -> tuple[_Array2T0, _Array2T1]: ...
+@overload
+def atleast_2d(a0: _Array2T, a1: _Array2T, /, *arys: _Array2T) -> tuple[_Array2T, ...]: ...  # type: ignore[overload-overlap]
+@overload
+def atleast_2d(a0: _ArrayLike[_ScalarT], /) -> _nt.MArray[_ScalarT]: ...
+@overload
+def atleast_2d(a0: _ArrayLike[_SCT0], a2: _ArrayLike[_SCT1], /) -> tuple[_nt.MArray[_SCT0], _nt.MArray[_SCT1]]: ...
+@overload
+def atleast_2d(
+    a0: _ArrayLike[_ScalarT], a2: _ArrayLike[_ScalarT], /, *arys: _ArrayLike[_ScalarT]
+) -> tuple[_nt.MArray[_ScalarT], ...]: ...
+@overload
+def atleast_2d(a0: ArrayLike, /) -> _nt.MArray[Any]: ...
+@overload
+def atleast_2d(a0: ArrayLike, a2: ArrayLike, /) -> tuple[_nt.MArray[Any], _nt.MArray[Any]]: ...
+@overload
+def atleast_2d(a0: ArrayLike, a2: ArrayLike, /, *arys: ArrayLike) -> tuple[_nt.MArray[Any], ...]: ...
+
+# keep in sync with `numpy._core.shape_base.atleast_2d`
+@overload
+def atleast_3d(a0: _Array3T, /) -> _Array3T: ...
+@overload
+def atleast_3d(a0: _Array3T0, a1: _Array3T1, /) -> tuple[_Array3T0, _Array3T1]: ...
+@overload
+def atleast_3d(a0: _Array3T, a1: _Array3T, /, *arys: _Array3T) -> tuple[_Array3T, ...]: ...  # type: ignore[overload-overlap]
+@overload
+def atleast_3d(a0: _ArrayLike[_ScalarT], /) -> _nt.MArray[_ScalarT]: ...
+@overload
+def atleast_3d(a0: _ArrayLike[_SCT0], a2: _ArrayLike[_SCT1], /) -> tuple[_nt.MArray[_SCT0], _nt.MArray[_SCT1]]: ...
+@overload
+def atleast_3d(
+    a0: _ArrayLike[_ScalarT], a2: _ArrayLike[_ScalarT], /, *arys: _ArrayLike[_ScalarT]
+) -> tuple[_nt.MArray[_ScalarT], ...]: ...
+@overload
+def atleast_3d(a0: ArrayLike, /) -> _nt.MArray[Any]: ...
+@overload
+def atleast_3d(a0: ArrayLike, a2: ArrayLike, /) -> tuple[_nt.MArray[Any], _nt.MArray[Any]]: ...
+@overload
+def atleast_3d(a0: ArrayLike, a2: ArrayLike, /, *arys: ArrayLike) -> tuple[_nt.MArray[Any], ...]: ...
+
+# keep in sync with `numpy._core.shape_base.vstack`
+@overload
+def vstack(
+    tup: Sequence[_ArrayLike[_ScalarT]], *, dtype: None = None, casting: np._CastingKind = "same_kind"
+) -> _nt.MArray[_ScalarT]: ...
+@overload
+def vstack(
+    tup: Sequence[ArrayLike], *, dtype: _DTypeLike[_ScalarT], casting: np._CastingKind = "same_kind"
+) -> _nt.MArray[_ScalarT]: ...
+@overload
+def vstack(
+    tup: Sequence[ArrayLike], *, dtype: DTypeLike | None = None, casting: np._CastingKind = "same_kind"
+) -> _nt.MArray[Any]: ...
+
+row_stack = vstack
+
+# keep in sync with `numpy._core.shape_base.hstack`
+@overload
+def hstack(
+    tup: Sequence[_ArrayLike[_ScalarT]], *, dtype: None = None, casting: np._CastingKind = "same_kind"
+) -> _nt.MArray[_ScalarT]: ...
+@overload
+def hstack(
+    tup: Sequence[ArrayLike], *, dtype: _DTypeLike[_ScalarT], casting: np._CastingKind = "same_kind"
+) -> _nt.MArray[_ScalarT]: ...
+@overload
+def hstack(
+    tup: Sequence[ArrayLike], *, dtype: DTypeLike | None = None, casting: np._CastingKind = "same_kind"
+) -> _nt.MArray[Any]: ...
+
+# keep in sync with `numpy._core.shape_base_impl.column_stack`
+@overload
+def column_stack(tup: Sequence[_ArrayLike[_ScalarT]]) -> _nt.MArray[_ScalarT]: ...
+@overload
+def column_stack(tup: Sequence[ArrayLike]) -> _nt.MArray[Incomplete]: ...
+
+# keep in sync with `numpy._core.shape_base_impl.dstack`
+@overload
+def dstack(tup: Sequence[_ArrayLike[_ScalarT]]) -> _nt.MArray[_ScalarT]: ...
+@overload
+def dstack(tup: Sequence[ArrayLike]) -> _nt.MArray[Incomplete]: ...
 
-vstack: _fromnxfunction_seq
-row_stack: _fromnxfunction_seq
-hstack: _fromnxfunction_seq
-column_stack: _fromnxfunction_seq
-dstack: _fromnxfunction_seq
-stack: _fromnxfunction_seq
+# keep in sync with `numpy._core.shape_base.stack`
+@overload
+def stack(
+    arrays: Sequence[_ArrayLike[_ScalarT]],
+    axis: SupportsIndex = 0,
+    out: None = None,
+    *,
+    dtype: None = None,
+    casting: np._CastingKind = "same_kind",
+) -> _nt.Array[_ScalarT]: ...
+@overload
+def stack(
+    arrays: Sequence[ArrayLike],
+    axis: SupportsIndex = 0,
+    out: None = None,
+    *,
+    dtype: _DTypeLike[_ScalarT],
+    casting: np._CastingKind = "same_kind",
+) -> _nt.Array[_ScalarT]: ...
+@overload
+def stack(
+    arrays: Sequence[ArrayLike],
+    axis: SupportsIndex = 0,
+    out: None = None,
+    *,
+    dtype: DTypeLike | None = None,
+    casting: np._CastingKind = "same_kind",
+) -> _nt.Array[Any]: ...
+@overload
+def stack(
+    arrays: Sequence[ArrayLike],
+    axis: SupportsIndex,
+    out: _ArrayT,
+    *,
+    dtype: DTypeLike | None = None,
+    casting: np._CastingKind = "same_kind",
+) -> _ArrayT: ...
+@overload
+def stack(
+    arrays: Sequence[ArrayLike],
+    axis: SupportsIndex = 0,
+    *,
+    out: _ArrayT,
+    dtype: DTypeLike | None = None,
+    casting: np._CastingKind = "same_kind",
+) -> _ArrayT: ...
 
-hsplit: _fromnxfunction_single
-diagflat: _fromnxfunction_single
+# keep in sync with `numpy._core.shape_base_impl.hsplit`
+@overload
+def hsplit(ary: _ArrayLike[_ScalarT], indices_or_sections: _ShapeLike) -> list[_nt.MArray[_ScalarT]]: ...
+@overload
+def hsplit(ary: ArrayLike, indices_or_sections: _ShapeLike) -> list[_nt.MArray[Incomplete]]: ...
 
+# keep in sync with `numpy._core.twodim_base_impl.hsplit`
+@overload
+def diagflat(v: _ArrayLike[_ScalarT], k: int = 0) -> _nt.MArray[_ScalarT]: ...
+@overload
+def diagflat(v: ArrayLike, k: int = 0) -> _nt.MArray[Incomplete]: ...
 def mask_rowcols(a: Incomplete, axis: Incomplete | None = None) -> _nt.MArray[Incomplete]: ...

tool/allowlists/todo.txt

@@ -6,11 +6,7 @@ numpy.ma.arange
 numpy.ma.argmax
 numpy.ma.argmin
 numpy.ma.asanyarray
-numpy.ma.atleast_1d
-numpy.ma.atleast_2d
-numpy.ma.atleast_3d
 numpy.ma.clip
-numpy.ma.column_stack
 numpy.ma.compress
 numpy.ma.copy
 numpy.ma.core._convert2ma
@@ -55,35 +51,15 @@ numpy.ma.core.trace
 numpy.ma.core.var
 numpy.ma.core.zeros
 numpy.ma.core.zeros_like
-numpy.ma.corrcoef
 numpy.ma.count
 numpy.ma.cumprod
 numpy.ma.cumsum
-numpy.ma.diagflat
 numpy.ma.diagonal
-numpy.ma.dstack
 numpy.ma.empty
 numpy.ma.empty_like
-numpy.ma.extras._fromnxfunction_allargs
-numpy.ma.extras._fromnxfunction_seq
-numpy.ma.extras._fromnxfunction_single
-numpy.ma.extras.atleast_1d
-numpy.ma.extras.atleast_2d
-numpy.ma.extras.atleast_3d
-numpy.ma.extras.column_stack
-numpy.ma.extras.corrcoef
-numpy.ma.extras.diagflat
-numpy.ma.extras.dstack
-numpy.ma.extras.hsplit
-numpy.ma.extras.hstack
-numpy.ma.extras.row_stack
-numpy.ma.extras.stack
-numpy.ma.extras.vstack
 numpy.ma.frombuffer
 numpy.ma.fromfunction
 numpy.ma.harden_mask
-numpy.ma.hsplit
-numpy.ma.hstack
 numpy.ma.identity
 numpy.ma.ids
 numpy.ma.indices
@@ -96,15 +72,12 @@ numpy.ma.prod
 numpy.ma.product
 numpy.ma.ravel
 numpy.ma.repeat
-numpy.ma.row_stack
 numpy.ma.soften_mask
 numpy.ma.squeeze
-numpy.ma.stack
 numpy.ma.std
 numpy.ma.sum
 numpy.ma.swapaxes
 numpy.ma.trace
 numpy.ma.var
-numpy.ma.vstack
 numpy.ma.zeros
 numpy.ma.zeros_like