From 325cf8113faaa738f5e376d9afd8a297513c4227 Mon Sep 17 00:00:00 2001 From: jorenham Date: Sun, 28 Dec 2025 22:57:34 +0100 Subject: [PATCH 1/8] =?UTF-8?q?=F0=9F=91=BD=EF=B8=8F=20`count=5Fnonzero`?= =?UTF-8?q?=20sync=20with=20upstream?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/_core/numeric.pyi | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/numpy-stubs/_core/numeric.pyi b/src/numpy-stubs/_core/numeric.pyi index e1e39dce..c556b77b 100644 --- a/src/numpy-stubs/_core/numeric.pyi +++ b/src/numpy-stubs/_core/numeric.pyi @@ -2030,6 +2030,10 @@ def full_like( @overload def count_nonzero(a: ArrayLike, axis: None = None, *, keepdims: L[False] = False) -> np.intp: ... @overload +def count_nonzero(a: _nt.ToGeneric_0d, axis: _ShapeLike | None = None, *, keepdims: L[True]) -> np.intp: ... +@overload +def count_nonzero(a: _nt.ToGeneric_1nd, axis: _ShapeLike | None = None, *, keepdims: L[True]) -> _nt.Array[np.intp]: ... +@overload def count_nonzero(a: ArrayLike, axis: _ShapeLike | None = None, *, keepdims: py_bool = False) -> Incomplete: ... # From f0b7e7084b06026dcd699ab81e22abfac9025c7c Mon Sep 17 00:00:00 2001 From: jorenham Date: Sun, 28 Dec 2025 22:57:44 +0100 Subject: [PATCH 2/8] =?UTF-8?q?=F0=9F=91=BD=EF=B8=8F=20`argwhere`=20sync?= =?UTF-8?q?=20with=20upstream?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/_core/numeric.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/numpy-stubs/_core/numeric.pyi b/src/numpy-stubs/_core/numeric.pyi index c556b77b..10dd61c6 100644 --- a/src/numpy-stubs/_core/numeric.pyi +++ b/src/numpy-stubs/_core/numeric.pyi @@ -2038,7 +2038,7 @@ def count_nonzero(a: ArrayLike, axis: _ShapeLike | None = None, *, keepdims: py_ # def flatnonzero(a: ArrayLike) -> _nt.Array1D[np.intp]: ... -def argwhere(a: ArrayLike) -> _nt.Array[np.intp]: ... +def argwhere(a: ArrayLike) -> _nt.Array2D[np.intp]: ... # def isfortran(a: _nt.Array | np.generic) -> py_bool: ... From ad8a632da540cc294523f8407c9190c60ffa51d9 Mon Sep 17 00:00:00 2001 From: jorenham Date: Sun, 28 Dec 2025 23:00:34 +0100 Subject: [PATCH 3/8] =?UTF-8?q?=F0=9F=91=BD=EF=B8=8F=20`indices`=20sync=20?= =?UTF-8?q?with=20upstream?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/_core/numeric.pyi | 82 ++++++++++++++++++++----------- 1 file changed, 54 insertions(+), 28 deletions(-) diff --git a/src/numpy-stubs/_core/numeric.pyi b/src/numpy-stubs/_core/numeric.pyi index 10dd61c6..de782a40 100644 --- a/src/numpy-stubs/_core/numeric.pyi +++ b/src/numpy-stubs/_core/numeric.pyi @@ -2075,7 +2075,7 @@ def correlate( a: _nt.CoComplex_1d | _nt.CoTimeDelta_1d | _nt.ToObject_1d, v: _nt.CoComplex_1d | _nt.CoTimeDelta_1d | _nt.ToObject_1d, mode: _Mode = "valid", -) -> _nt.Array1D[Any]: ... +) -> _nt.Array1D: ... # @overload @@ -2107,7 +2107,7 @@ def convolve( a: _nt.CoComplex_1d | _nt.CoTimeDelta_1d | _nt.ToObject_1d, v: _nt.CoComplex_1d | _nt.CoTimeDelta_1d | _nt.ToObject_1d, mode: _Mode = "full", -) -> _nt.Array1D[Any]: ... +) -> _nt.Array1D: ... # @overload @@ -2271,42 +2271,68 @@ def cross( ) -> _nt.Array[Any]: ... # -@overload +@overload # 0d, dtype=int (default), sparse=False (default) +def indices(dimensions: tuple[()], dtype: type[int] = int, sparse: L[False] = False) -> _nt.Array1D[np.intp]: ... # noqa: PYI011 +@overload # 0d, dtype=, sparse=True +def indices(dimensions: tuple[()], dtype: DTypeLike | None = int, *, sparse: L[True]) -> tuple[()]: ... # noqa: PYI011 +@overload # 0d, dtype=, sparse=False (default) +def indices(dimensions: tuple[()], dtype: _DTypeLike[_ScalarT], sparse: L[False] = False) -> _nt.Array1D[_ScalarT]: ... +@overload # 0d, dtype=, sparse=False (default) +def indices(dimensions: tuple[()], dtype: DTypeLike, sparse: L[False] = False) -> _nt.Array1D: ... +@overload # 1d, dtype=int (default), sparse=False (default) +def indices(dimensions: tuple[int], dtype: type[int] = int, sparse: L[False] = False) -> _nt.Array2D[np.intp]: ... # noqa: PYI011 +@overload # 1d, dtype=int (default), sparse=True +def indices(dimensions: tuple[int], dtype: type[int] = int, *, sparse: L[True]) -> tuple[_nt.Array1D[np.intp]]: ... # noqa: PYI011 +@overload # 1d, dtype=, sparse=False (default) +def indices(dimensions: tuple[int], dtype: _DTypeLike[_ScalarT], sparse: L[False] = False) -> _nt.Array2D[_ScalarT]: ... +@overload # 1d, dtype=, sparse=True +def indices(dimensions: tuple[int], dtype: _DTypeLike[_ScalarT], sparse: L[True]) -> tuple[_nt.Array1D[_ScalarT]]: ... +@overload # 1d, dtype=, sparse=False (default) +def indices(dimensions: tuple[int], dtype: DTypeLike, sparse: L[False] = False) -> _nt.Array2D: ... +@overload # 1d, dtype=, sparse=True +def indices(dimensions: tuple[int], dtype: DTypeLike, sparse: L[True]) -> tuple[_nt.Array1D]: ... +@overload # 2d, dtype=int (default), sparse=False (default) +def indices(dimensions: tuple[int, int], dtype: type[int] = int, sparse: L[False] = False) -> _nt.Array3D[np.intp]: ... # noqa: PYI011 +@overload # 2d, dtype=int (default), sparse=True def indices( - dimensions: _nt.ToInteger_1d, - dtype: type[_nt.JustInt] = int, # noqa: PYI011 - sparse: L[False] = False, -) -> _nt.Array[np.intp]: ... -@overload + dimensions: tuple[int, int], + dtype: type[int] = int, # noqa: PYI011 + *, + sparse: L[True], +) -> tuple[_nt.Array2D[np.intp], _nt.Array2D[np.intp]]: ... +@overload # 2d, dtype=, sparse=False (default) def indices( - dimensions: _nt.ToInteger_1d, dtype: type[_nt.JustInt], sparse: L[True] -) -> tuple[_nt.Array[np.intp], ...]: ... -@overload + dimensions: tuple[int, int], dtype: _DTypeLike[_ScalarT], sparse: L[False] = False +) -> _nt.Array3D[_ScalarT]: ... +@overload # 2d, dtype=, sparse=True def indices( - dimensions: _nt.ToInteger_1d, - dtype: type[_nt.JustInt] = int, # noqa: PYI011 + dimensions: tuple[int, int], dtype: _DTypeLike[_ScalarT], sparse: L[True] +) -> tuple[_nt.Array2D[_ScalarT], _nt.Array2D[_ScalarT]]: ... +@overload # 2d, dtype=, sparse=False (default) +def indices(dimensions: tuple[int, int], dtype: DTypeLike, sparse: L[False] = False) -> _nt.Array3D: ... +@overload # 2d, dtype=, sparse=True +def indices(dimensions: tuple[int, int], dtype: DTypeLike, sparse: L[True]) -> tuple[_nt.Array2D, _nt.Array2D]: ... +@overload # ?d, dtype=int (default), sparse=False (default) +def indices(dimensions: Sequence[int], dtype: type[int] = int, sparse: L[False] = False) -> _nt.Array[np.intp]: ... # noqa: PYI011 +@overload # ?d, dtype=int (default), sparse=True +def indices( + dimensions: Sequence[int], + dtype: type[int] = int, # noqa: PYI011 *, sparse: L[True], ) -> tuple[_nt.Array[np.intp], ...]: ... -@overload +@overload # ?d, dtype=, sparse=False (default) def indices( - dimensions: _nt.ToInteger_1d, dtype: _DTypeLike[_ScalarT], sparse: L[False] = False + dimensions: Sequence[int], dtype: _DTypeLike[_ScalarT], sparse: L[False] = False ) -> _nt.Array[_ScalarT]: ... -@overload +@overload # ?d, dtype=, sparse=True def indices( - dimensions: _nt.ToInteger_1d, dtype: _DTypeLike[_ScalarT], sparse: L[True] + dimensions: Sequence[int], dtype: _DTypeLike[_ScalarT], sparse: L[True] ) -> tuple[_nt.Array[_ScalarT], ...]: ... -@overload -def indices(dimensions: _nt.ToInteger_1d, dtype: DTypeLike | None = int, sparse: L[False] = False) -> _nt.Array: ... # noqa: PYI011 -@overload -def indices(dimensions: _nt.ToInteger_1d, dtype: DTypeLike | None, sparse: L[True]) -> tuple[_nt.Array, ...]: ... -@overload -def indices( - dimensions: _nt.ToInteger_1d, - dtype: DTypeLike | None = int, # noqa: PYI011 - *, - sparse: L[True], -) -> tuple[_nt.Array, ...]: ... +@overload # ?d, dtype=, sparse=False (default) +def indices(dimensions: Sequence[int], dtype: DTypeLike, sparse: L[False] = False) -> ndarray: ... +@overload # ?d, dtype=, sparse=True +def indices(dimensions: Sequence[int], dtype: DTypeLike, sparse: L[True]) -> tuple[ndarray, ...]: ... # keep in sync with `ma.core.fromfunction` def fromfunction( From 3970b00fbf5b71bb73926c39e34835a03bcb1f16 Mon Sep 17 00:00:00 2001 From: jorenham Date: Sun, 28 Dec 2025 23:03:13 +0100 Subject: [PATCH 4/8] =?UTF-8?q?=F0=9F=91=BD=EF=B8=8F=20`identity`=20sync?= =?UTF-8?q?=20with=20upstream?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/_core/numeric.pyi | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/src/numpy-stubs/_core/numeric.pyi b/src/numpy-stubs/_core/numeric.pyi index de782a40..f8e99b56 100644 --- a/src/numpy-stubs/_core/numeric.pyi +++ b/src/numpy-stubs/_core/numeric.pyi @@ -2352,14 +2352,26 @@ def binary_repr(num: SupportsIndex, width: int | None = None) -> str: ... def base_repr(number: SupportsAbs[float], base: float = 2, padding: SupportsIndex = 0) -> str: ... # -@overload +@overload # dtype: None (default) def identity(n: int, dtype: None = None, *, like: _SupportsArrayFunc | None = None) -> _nt.Array2D[np.float64]: ... -@overload +@overload # dtype: known scalar type def identity( n: int, dtype: _DTypeLike[_ScalarT], *, like: _SupportsArrayFunc | None = None ) -> _nt.Array2D[_ScalarT]: ... -@overload -def identity(n: int, dtype: DTypeLike | None = None, *, like: _SupportsArrayFunc | None = None) -> _nt.Array2D: ... +@overload # dtype: like bool +def identity(n: int, dtype: _nt.ToDTypeBool, *, like: _SupportsArrayFunc | None = None) -> _nt.Array2D[np.bool]: ... +@overload # dtype: like int_ +def identity(n: int, dtype: _nt.ToDTypeInt64, *, like: _SupportsArrayFunc | None = None) -> _nt.Array2D[np.int64]: ... +@overload # dtype: like float64 +def identity( + n: int, dtype: _nt.ToDTypeFloat64, *, like: _SupportsArrayFunc | None = None +) -> _nt.Array2D[np.float64]: ... +@overload # dtype: like complex128 +def identity( + n: int, dtype: _nt.ToDTypeComplex128, *, like: _SupportsArrayFunc | None = None +) -> _nt.Array2D[np.complex128]: ... +@overload # dtype: unknown +def identity(n: int, dtype: DTypeLike, *, like: _SupportsArrayFunc | None = None) -> _nt.Array2D[Incomplete]: ... # def allclose( From 19f9cab3385e4764608199b33f4018b877c62bcb Mon Sep 17 00:00:00 2001 From: jorenham Date: Sun, 28 Dec 2025 23:08:56 +0100 Subject: [PATCH 5/8] =?UTF-8?q?=F0=9F=91=BD=EF=B8=8F=20`isclose`=20sync=20?= =?UTF-8?q?with=20upstream?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/_core/numeric.pyi | 56 ++++++++++++++++++++++++------- 1 file changed, 44 insertions(+), 12 deletions(-) diff --git a/src/numpy-stubs/_core/numeric.pyi b/src/numpy-stubs/_core/numeric.pyi index f8e99b56..a842a713 100644 --- a/src/numpy-stubs/_core/numeric.pyi +++ b/src/numpy-stubs/_core/numeric.pyi @@ -2379,30 +2379,62 @@ def allclose( ) -> py_bool: ... # -@overload +@overload # scalar, scalar +def isclose( + a: _nt.CoComplex_0d, b: _nt.CoComplex_0d, rtol: ArrayLike = 1e-5, atol: ArrayLike = 1e-8, equal_nan: py_bool = False +) -> np.bool_: ... +@overload # known shape, same shape or scalar def isclose( - a: np.generic | _nt._PyScalar, - b: np.generic | _nt._PyScalar, + a: np.ndarray[_ShapeT], + b: np.ndarray[_ShapeT] | _nt.CoComplex_0d, rtol: ArrayLike = 1e-5, atol: ArrayLike = 1e-8, equal_nan: py_bool = False, -) -> np.bool: ... -@overload +) -> _nt.Array[np.bool_, _ShapeT]: ... +@overload # same shape or scalar, known shape def isclose( - a: _nt.ToGeneric_1nd, - b: _nt.ToGeneric_nd, + a: np.ndarray[_ShapeT] | _nt.CoComplex_0d, + b: np.ndarray[_ShapeT], rtol: ArrayLike = 1e-5, atol: ArrayLike = 1e-8, equal_nan: py_bool = False, -) -> _nt.Array[np.bool]: ... -@overload +) -> _nt.Array[np.bool_, _ShapeT]: ... +@overload # 1d sequence, <=1d array-like def isclose( - a: _nt.ToGeneric_nd, - b: _nt.ToGeneric_1nd, + a: _nt.CoComplex_1ds, + b: _nt.CoComplex_1ds | _nt.CoComplex_0d, rtol: ArrayLike = 1e-5, atol: ArrayLike = 1e-8, equal_nan: py_bool = False, -) -> _nt.Array[np.bool]: ... +) -> _nt.Array1D[np.bool_]: ... +@overload # <=1d array-like, 1d sequence +def isclose( + a: _nt.CoComplex_1ds | _nt.CoComplex_0d, + b: _nt.CoComplex_1ds, + rtol: ArrayLike = 1e-5, + atol: ArrayLike = 1e-8, + equal_nan: py_bool = False, +) -> _nt.Array1D[np.bool_]: ... +@overload # 2d sequence, <=2d array-like +def isclose( + a: _nt.CoComplex_2ds, + b: _nt.CoComplex_2ds | _nt.CoComplex_1ds | _nt.CoComplex_0d, + rtol: ArrayLike = 1e-5, + atol: ArrayLike = 1e-8, + equal_nan: py_bool = False, +) -> _nt.Array1D[np.bool_]: ... +@overload # <=2d array-like, 2d sequence +def isclose( + b: _nt.CoComplex_2ds | _nt.CoComplex_1ds | _nt.CoComplex_0d, + a: _nt.CoComplex_2ds, + rtol: ArrayLike = 1e-5, + atol: ArrayLike = 1e-8, + equal_nan: py_bool = False, +) -> _nt.Array1D[np.bool_]: ... +@overload # unknown shape, unknown shape +def isclose( + a: ArrayLike, b: ArrayLike, rtol: ArrayLike = 1e-5, atol: ArrayLike = 1e-8, equal_nan: py_bool = False +) -> _nt.Array[np.bool_] | Any: ... # def array_equal(a1: ArrayLike, a2: ArrayLike, equal_nan: py_bool = False) -> py_bool: ... From 4a19e660c7c0e8d52c1ba92fb6aefccdd64db0f0 Mon Sep 17 00:00:00 2001 From: jorenham Date: Sun, 28 Dec 2025 23:10:51 +0100 Subject: [PATCH 6/8] =?UTF-8?q?=F0=9F=8F=B7=EF=B8=8F=20prefer=20`=5Fnumtyp?= =?UTF-8?q?e.Array`=20over=20a=20bare=20`ndarray`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/_core/numeric.pyi | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/numpy-stubs/_core/numeric.pyi b/src/numpy-stubs/_core/numeric.pyi index a842a713..d018ada4 100644 --- a/src/numpy-stubs/_core/numeric.pyi +++ b/src/numpy-stubs/_core/numeric.pyi @@ -2444,8 +2444,8 @@ def array_equiv(a1: ArrayLike, a2: ArrayLike) -> py_bool: ... @overload def astype( x: _nt.Array[Any, _ShapeT], dtype: _DTypeLike[_ScalarT], /, *, copy: py_bool = True, device: _Device | None = None -) -> ndarray[_ShapeT, dtype[_ScalarT]]: ... +) -> _nt.Array[_ScalarT, _ShapeT]: ... @overload def astype( x: _nt.Array[Any, _ShapeT], dtype: DTypeLike | None, /, *, copy: py_bool = True, device: _Device | None = None -) -> ndarray[_ShapeT, dtype]: ... +) -> _nt.Array[Incomplete, _ShapeT]: ... From 81979262315afa4bfacf415bfd5ab1c77df7303d Mon Sep 17 00:00:00 2001 From: jorenham Date: Sun, 28 Dec 2025 23:27:26 +0100 Subject: [PATCH 7/8] =?UTF-8?q?=F0=9F=A9=B9=20fix=20`isclose`=20shape-typi?= =?UTF-8?q?ng=20issue?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/_core/numeric.pyi | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/numpy-stubs/_core/numeric.pyi b/src/numpy-stubs/_core/numeric.pyi index d018ada4..30abfb7b 100644 --- a/src/numpy-stubs/_core/numeric.pyi +++ b/src/numpy-stubs/_core/numeric.pyi @@ -2381,7 +2381,11 @@ def allclose( # @overload # scalar, scalar def isclose( - a: _nt.CoComplex_0d, b: _nt.CoComplex_0d, rtol: ArrayLike = 1e-5, atol: ArrayLike = 1e-8, equal_nan: py_bool = False + a: _nt.co_complex | complex, + b: _nt.co_complex | complex, + rtol: ArrayLike = 1e-5, + atol: ArrayLike = 1e-8, + equal_nan: py_bool = False, ) -> np.bool_: ... @overload # known shape, same shape or scalar def isclose( From 7df277ced2633e022a684873db088a13aeef06b3 Mon Sep 17 00:00:00 2001 From: jorenham Date: Sun, 28 Dec 2025 23:28:23 +0100 Subject: [PATCH 8/8] =?UTF-8?q?=E2=9C=85=20update=20`=5Fcore.numeric`=20ty?= =?UTF-8?q?pe-tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../@test/static/accept/array_constructors.pyi | 2 +- src/numpy-stubs/@test/static/accept/numeric.pyi | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/numpy-stubs/@test/static/accept/array_constructors.pyi b/src/numpy-stubs/@test/static/accept/array_constructors.pyi index b43bce04..e80d2d70 100644 --- a/src/numpy-stubs/@test/static/accept/array_constructors.pyi +++ b/src/numpy-stubs/@test/static/accept/array_constructors.pyi @@ -263,7 +263,7 @@ assert_type(np.fromfunction(func, (3, 5)), MyArray[np.float64]) assert_type(np.identity(10), _nt.Array2D[np.float64]) assert_type(np.identity(10, dtype=np.int64), _nt.Array2D[np.int64]) -assert_type(np.identity(10, dtype=int), _nt.Array2D) +assert_type(np.identity(10, dtype=int), _nt.Array2D[np.int64]) assert_type(np.atleast_1d(A), _nt.Array[np.float64]) assert_type(np.atleast_1d(C), _nt.Array) diff --git a/src/numpy-stubs/@test/static/accept/numeric.pyi b/src/numpy-stubs/@test/static/accept/numeric.pyi index 94ca4187..6da05f0b 100644 --- a/src/numpy-stubs/@test/static/accept/numeric.pyi +++ b/src/numpy-stubs/@test/static/accept/numeric.pyi @@ -25,14 +25,14 @@ AR_O: _nt.Array[np.object_] assert_type(np.count_nonzero(i8), np.intp) assert_type(np.count_nonzero(AR_i8), np.intp) assert_type(np.count_nonzero(ints), np.intp) -assert_type(np.count_nonzero(AR_i8, keepdims=True), Any) +assert_type(np.count_nonzero(AR_i8, keepdims=True), np.intp) assert_type(np.count_nonzero(AR_i8, axis=0), Any) assert_type(np.isfortran(i8), bool) assert_type(np.isfortran(AR_i8), bool) -assert_type(np.argwhere(i8), _nt.Array[np.intp]) -assert_type(np.argwhere(AR_i8), _nt.Array[np.intp]) +assert_type(np.argwhere(i8), _nt.Array2D[np.intp]) +assert_type(np.argwhere(AR_i8), _nt.Array2D[np.intp]) assert_type(np.flatnonzero(i8), _nt.Array1D[np.intp]) assert_type(np.flatnonzero(AR_i8), _nt.Array1D[np.intp]) @@ -121,7 +121,7 @@ assert_type(np.allclose(AR_i8, AR_i8), bool) assert_type(np.isclose(i8, i8), np.bool) assert_type(np.isclose(i8, AR_i8), _nt.Array[np.bool]) -assert_type(np.isclose(ints, AR_i8), _nt.Array[np.bool]) +assert_type(np.isclose(ints, AR_i8), _nt.Array1D[np.bool]) assert_type(np.isclose(AR_i8, AR_i8), _nt.Array[np.bool]) assert_type(np.array_equal(i8, AR_i8), bool)