From 6ee7170c107e39a80295554e6fe5b189adb25f18 Mon Sep 17 00:00:00 2001 From: jorenham Date: Tue, 23 Dec 2025 20:16:10 +0100 Subject: [PATCH 1/2] =?UTF-8?q?=F0=9F=91=BD=EF=B8=8F=20`ma.extras`=20sync?= =?UTF-8?q?=20the=20`nxfunction`=20stuff=20with=20upstream?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/ma/extras.pyi | 215 +++++++++++++++++++++++++++++----- tool/allowlists/todo.txt | 25 ---- 2 files changed, 184 insertions(+), 56 deletions(-) diff --git a/src/numpy-stubs/ma/extras.pyi b/src/numpy-stubs/ma/extras.pyi index cc5e9d44..42816c10 100644 --- a/src/numpy-stubs/ma/extras.pyi +++ b/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__ = () @@ -161,18 +162,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]: ... diff --git a/tool/allowlists/todo.txt b/tool/allowlists/todo.txt index 23d65d2c..79650681 100644 --- a/tool/allowlists/todo.txt +++ b/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 @@ -61,29 +57,11 @@ 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 +74,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 From 8698deda80d6bf16d8853be22a1a26ff2dfa3174 Mon Sep 17 00:00:00 2001 From: jorenham Date: Tue, 23 Dec 2025 20:17:34 +0100 Subject: [PATCH 2/2] =?UTF-8?q?=F0=9F=91=BD=EF=B8=8F=20`ma.corrcoef`=20syn?= =?UTF-8?q?c=20the=20with=20upstream?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/ma/extras.pyi | 9 +-------- tool/allowlists/todo.txt | 2 -- 2 files changed, 1 insertion(+), 10 deletions(-) diff --git a/src/numpy-stubs/ma/extras.pyi b/src/numpy-stubs/ma/extras.pyi index 42816c10..cc102f25 100644 --- a/src/numpy-stubs/ma/extras.pyi +++ b/src/numpy-stubs/ma/extras.pyi @@ -134,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: ... diff --git a/tool/allowlists/todo.txt b/tool/allowlists/todo.txt index 79650681..af32d2ea 100644 --- a/tool/allowlists/todo.txt +++ b/tool/allowlists/todo.txt @@ -51,11 +51,9 @@ 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.empty numpy.ma.empty_like