From 5dec9a612fadaf640298caa7047209561acc50c7 Mon Sep 17 00:00:00 2001 From: jorenham Date: Sat, 27 Dec 2025 13:54:15 +0100 Subject: [PATCH 1/9] =?UTF-8?q?=E2=9C=A8=20`histogram2d`=20shape-typing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/lib/_twodim_base_impl.pyi | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/numpy-stubs/lib/_twodim_base_impl.pyi b/src/numpy-stubs/lib/_twodim_base_impl.pyi index 51cc4301..abff9c37 100644 --- a/src/numpy-stubs/lib/_twodim_base_impl.pyi +++ b/src/numpy-stubs/lib/_twodim_base_impl.pyi @@ -38,8 +38,8 @@ _MaskFunc: TypeAlias = Callable[ ] _Sequence01D: TypeAlias = _T | Sequence[_T] -_Indices2D: TypeAlias = tuple[_nt.Array[np.intp], _nt.Array[np.intp]] -_Histogram2D: TypeAlias = tuple[_nt.Array[np.float64], _nt.Array[_ScalarT], _nt.Array[_ScalarT]] +_Indices2D: TypeAlias = tuple[_nt.Array1D[np.intp], _nt.Array1D[np.intp]] +_Histogram2D: TypeAlias = tuple[_nt.Array2D[np.float64], _nt.Array1D[_ScalarT], _nt.Array1D[_ScalarT]] ### From 3f59cca4fb42f2df4bf062aaba24d56755195e5d Mon Sep 17 00:00:00 2001 From: jorenham Date: Sat, 27 Dec 2025 13:55:22 +0100 Subject: [PATCH 2/9] =?UTF-8?q?=E2=9C=A8=20`fliplr`=20and=20`flipud`=20sha?= =?UTF-8?q?pe-typing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/lib/_twodim_base_impl.pyi | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/numpy-stubs/lib/_twodim_base_impl.pyi b/src/numpy-stubs/lib/_twodim_base_impl.pyi index abff9c37..cc21425a 100644 --- a/src/numpy-stubs/lib/_twodim_base_impl.pyi +++ b/src/numpy-stubs/lib/_twodim_base_impl.pyi @@ -30,6 +30,7 @@ _ScalarT = TypeVar("_ScalarT", bound=np.generic) _ComplexT = TypeVar("_ComplexT", bound=np.complexfloating) _InexactT = TypeVar("_InexactT", bound=np.inexact) _CoComplexT = TypeVar("_CoComplexT", bound=_nt.co_complex) +_ArrayT = TypeVar("_ArrayT", bound=_nt.Array) # The returned arrays dtype must be compatible with `np.equal` _Device: TypeAlias = L["cpu"] @@ -43,6 +44,8 @@ _Histogram2D: TypeAlias = tuple[_nt.Array2D[np.float64], _nt.Array1D[_ScalarT], ### +@overload +def fliplr(m: _ArrayT) -> _ArrayT: ... @overload def fliplr(m: _nt._ToArray_nd[_ScalarT]) -> _nt.Array[_ScalarT]: ... @overload @@ -50,6 +53,8 @@ def fliplr(m: _nt.ToGeneric_nd) -> _nt.Array: ... # @overload +def flipud(m: _ArrayT) -> _ArrayT: ... +@overload def flipud(m: _nt._ToArray_nd[_ScalarT]) -> _nt.Array[_ScalarT]: ... @overload def flipud(m: _nt.ToGeneric_nd) -> _nt.Array: ... From e5f52098a3fa9be69296549bf0fe76fc37a9bb76 Mon Sep 17 00:00:00 2001 From: jorenham Date: Sat, 27 Dec 2025 14:00:05 +0100 Subject: [PATCH 3/9] =?UTF-8?q?=E2=9C=A8=20`diag`=20shape-typing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/lib/_twodim_base_impl.pyi | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/src/numpy-stubs/lib/_twodim_base_impl.pyi b/src/numpy-stubs/lib/_twodim_base_impl.pyi index cc21425a..51b3f051 100644 --- a/src/numpy-stubs/lib/_twodim_base_impl.pyi +++ b/src/numpy-stubs/lib/_twodim_base_impl.pyi @@ -1,5 +1,6 @@ +from _typeshed import Incomplete from collections.abc import Callable, Sequence -from typing import Literal as L, TypeAlias, overload +from typing import Literal as L, Never, TypeAlias, overload from typing_extensions import TypeVar import _numtype as _nt @@ -32,6 +33,11 @@ _InexactT = TypeVar("_InexactT", bound=np.inexact) _CoComplexT = TypeVar("_CoComplexT", bound=_nt.co_complex) _ArrayT = TypeVar("_ArrayT", bound=_nt.Array) +# Workaround for mypy's and pyright's lack of compliance with the typing spec for +# overloads for gradual types. This works because only `Any` and `Never` are assignable +# to `Never`. +_ArrayNoD: TypeAlias = np.ndarray[tuple[Never] | tuple[Never, Never], np.dtype[_ScalarT]] + # The returned arrays dtype must be compatible with `np.equal` _Device: TypeAlias = L["cpu"] _MaskFunc: TypeAlias = Callable[ @@ -65,7 +71,7 @@ def eye( N: int, M: int | None = None, k: int = 0, - dtype: _nt.ToDTypeFloat64 | None = ..., + dtype: _nt.ToDTypeFloat64 | None = ..., # = float order: _OrderCF = "C", *, device: _Device | None = None, @@ -98,7 +104,7 @@ def eye( N: int, M: int | None = None, k: int = 0, - dtype: _nt.ToDType = ..., + dtype: _nt.ToDType = ..., # = float order: _OrderCF = "C", *, device: _Device | None = None, @@ -107,6 +113,16 @@ def eye( # @overload +def diag(v: _ArrayNoD[_ScalarT], k: int = 0) -> _nt.Array[_ScalarT]: ... # type: ignore[overload-overlap] +@overload +def diag(v: _nt.Array2D[_ScalarT] | Sequence[Sequence[_ScalarT]], k: int = 0) -> _nt.Array1D[_ScalarT]: ... +@overload +def diag(v: _nt.Array1D[_ScalarT] | Sequence[_ScalarT], k: int = 0) -> _nt.Array2D[_ScalarT]: ... +@overload +def diag(v: Sequence[Sequence[_nt.ToGeneric_0d]], k: int = 0) -> _nt.Array1D[Incomplete]: ... +@overload +def diag(v: Sequence[_nt.ToGeneric_0d], k: int = 0) -> _nt.Array2D[Incomplete]: ... +@overload def diag(v: _nt._ToArray_nd[_ScalarT], k: int = 0) -> _nt.Array[_ScalarT]: ... @overload def diag(v: _nt.ToGeneric_nd, k: int = 0) -> _nt.Array: ... From f1440b713d312ed62cd8facc5f919f96ff1467f8 Mon Sep 17 00:00:00 2001 From: jorenham Date: Sat, 27 Dec 2025 14:00:25 +0100 Subject: [PATCH 4/9] =?UTF-8?q?=E2=9C=A8=20`diagflat`=20shape-typing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/lib/_twodim_base_impl.pyi | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/numpy-stubs/lib/_twodim_base_impl.pyi b/src/numpy-stubs/lib/_twodim_base_impl.pyi index 51b3f051..be045f6d 100644 --- a/src/numpy-stubs/lib/_twodim_base_impl.pyi +++ b/src/numpy-stubs/lib/_twodim_base_impl.pyi @@ -129,9 +129,9 @@ def diag(v: _nt.ToGeneric_nd, k: int = 0) -> _nt.Array: ... # @overload -def diagflat(v: _nt._ToArray_nd[_ScalarT], k: int = 0) -> _nt.Array[_ScalarT]: ... +def diagflat(v: _nt._ToArray_nd[_ScalarT], k: int = 0) -> _nt.Array2D[_ScalarT]: ... @overload -def diagflat(v: _nt.ToGeneric_nd, k: int = 0) -> _nt.Array: ... +def diagflat(v: _nt.ToGeneric_nd, k: int = 0) -> _nt.Array2D[Incomplete]: ... # @overload From 2972bf0fa9ced9952fb515636e58a0acfe698a03 Mon Sep 17 00:00:00 2001 From: jorenham Date: Sat, 27 Dec 2025 14:01:19 +0100 Subject: [PATCH 5/9] =?UTF-8?q?=E2=9C=A8=20`tri`=20shape-typing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/lib/_twodim_base_impl.pyi | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/numpy-stubs/lib/_twodim_base_impl.pyi b/src/numpy-stubs/lib/_twodim_base_impl.pyi index be045f6d..8ee40425 100644 --- a/src/numpy-stubs/lib/_twodim_base_impl.pyi +++ b/src/numpy-stubs/lib/_twodim_base_impl.pyi @@ -137,19 +137,19 @@ def diagflat(v: _nt.ToGeneric_nd, k: int = 0) -> _nt.Array2D[Incomplete]: ... @overload def tri( N: int, M: int | None = None, k: int = 0, dtype: _nt.ToDTypeFloat64 | None = ..., *, like: _Like | None = None -) -> _nt.Array[np.float64]: ... +) -> _nt.Array2D[np.float64]: ... @overload def tri( N: int, M: int | None, k: int, dtype: _nt._ToDType[_ScalarT], *, like: _Like | None = None -) -> _nt.Array[_ScalarT]: ... +) -> _nt.Array2D[_ScalarT]: ... @overload def tri( N: int, M: int | None = None, k: int = 0, *, dtype: _nt._ToDType[_ScalarT], like: _Like | None = None -) -> _nt.Array[_ScalarT]: ... +) -> _nt.Array2D[_ScalarT]: ... @overload def tri( N: int, M: int | None = None, k: int = 0, dtype: _nt.ToDType = ..., *, like: _Like | None = None -) -> _nt.Array: ... +) -> _nt.Array2D: ... # @overload From 5fb4e5001a985366508c498d05b5d58352b2846a Mon Sep 17 00:00:00 2001 From: jorenham Date: Sat, 27 Dec 2025 14:04:03 +0100 Subject: [PATCH 6/9] =?UTF-8?q?=E2=9C=A8=20`vander`=20shape-typing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/lib/_twodim_base_impl.pyi | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/numpy-stubs/lib/_twodim_base_impl.pyi b/src/numpy-stubs/lib/_twodim_base_impl.pyi index 8ee40425..d01d4d53 100644 --- a/src/numpy-stubs/lib/_twodim_base_impl.pyi +++ b/src/numpy-stubs/lib/_twodim_base_impl.pyi @@ -31,6 +31,7 @@ _ScalarT = TypeVar("_ScalarT", bound=np.generic) _ComplexT = TypeVar("_ComplexT", bound=np.complexfloating) _InexactT = TypeVar("_InexactT", bound=np.inexact) _CoComplexT = TypeVar("_CoComplexT", bound=_nt.co_complex) +_NumberT = TypeVar("_NumberT", bound=np.number) _ArrayT = TypeVar("_ArrayT", bound=_nt.Array) # Workaround for mypy's and pyright's lack of compliance with the typing spec for @@ -189,15 +190,17 @@ def triu(m: _nt.ToGeneric_nd, k: int = 0) -> _nt.Array: ... # @overload -def vander(x: _nt.ToBool_1d, N: int | None = None, increasing: bool = False) -> _nt.Array[np.intp]: ... +def vander(x: _nt._ToArray_1d[_NumberT], N: int | None = None, increasing: bool = False) -> _nt.Array2D[_NumberT]: ... @overload -def vander(x: _nt.ToInteger_1d, N: int | None = None, increasing: bool = False) -> _nt.Array[np.signedinteger]: ... +def vander(x: _nt.ToBool_1d, N: int | None = None, increasing: bool = False) -> _nt.Array2D[np.intp]: ... @overload -def vander(x: _nt.ToFloating_1d, N: int | None = None, increasing: bool = False) -> _nt.Array[np.floating]: ... +def vander(x: _nt.ToInteger_1d, N: int | None = None, increasing: bool = False) -> _nt.Array2D[np.signedinteger]: ... @overload -def vander(x: _nt.ToComplex_1d, N: int | None = None, increasing: bool = False) -> _nt.Array[np.complexfloating]: ... +def vander(x: _nt.ToFloating_1d, N: int | None = None, increasing: bool = False) -> _nt.Array2D[np.floating]: ... @overload -def vander(x: _nt.ToObject_1d, N: int | None = None, increasing: bool = False) -> _nt.Array[np.object_]: ... +def vander(x: _nt.ToComplex_1d, N: int | None = None, increasing: bool = False) -> _nt.Array2D[np.complexfloating]: ... +@overload +def vander(x: _nt.ToObject_1d, N: int | None = None, increasing: bool = False) -> _nt.Array2D[np.object_]: ... # @overload From 8145ae6335acd8e69b072f4e99273f79ee78277e Mon Sep 17 00:00:00 2001 From: jorenham Date: Sat, 27 Dec 2025 14:14:02 +0100 Subject: [PATCH 7/9] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20`vander`=20narrow=20re?= =?UTF-8?q?turn=20dtype?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/lib/_twodim_base_impl.pyi | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/numpy-stubs/lib/_twodim_base_impl.pyi b/src/numpy-stubs/lib/_twodim_base_impl.pyi index d01d4d53..81351397 100644 --- a/src/numpy-stubs/lib/_twodim_base_impl.pyi +++ b/src/numpy-stubs/lib/_twodim_base_impl.pyi @@ -31,7 +31,6 @@ _ScalarT = TypeVar("_ScalarT", bound=np.generic) _ComplexT = TypeVar("_ComplexT", bound=np.complexfloating) _InexactT = TypeVar("_InexactT", bound=np.inexact) _CoComplexT = TypeVar("_CoComplexT", bound=_nt.co_complex) -_NumberT = TypeVar("_NumberT", bound=np.number) _ArrayT = TypeVar("_ArrayT", bound=_nt.Array) # Workaround for mypy's and pyright's lack of compliance with the typing spec for @@ -190,15 +189,13 @@ def triu(m: _nt.ToGeneric_nd, k: int = 0) -> _nt.Array: ... # @overload -def vander(x: _nt._ToArray_1d[_NumberT], N: int | None = None, increasing: bool = False) -> _nt.Array2D[_NumberT]: ... +def vander(x: _nt.CoInteger_1d, N: int | None = None, increasing: bool = False) -> _nt.Array2D[np.int_]: ... @overload -def vander(x: _nt.ToBool_1d, N: int | None = None, increasing: bool = False) -> _nt.Array2D[np.intp]: ... +def vander(x: _nt._ToArray_1d[_InexactT], N: int | None = None, increasing: bool = False) -> _nt.Array2D[_InexactT]: ... @overload -def vander(x: _nt.ToInteger_1d, N: int | None = None, increasing: bool = False) -> _nt.Array2D[np.signedinteger]: ... +def vander(x: _nt.ToFloat64_1d, N: int | None = None, increasing: bool = False) -> _nt.Array2D[np.float64]: ... @overload -def vander(x: _nt.ToFloating_1d, N: int | None = None, increasing: bool = False) -> _nt.Array2D[np.floating]: ... -@overload -def vander(x: _nt.ToComplex_1d, N: int | None = None, increasing: bool = False) -> _nt.Array2D[np.complexfloating]: ... +def vander(x: _nt.ToComplex128_1d, N: int | None = None, increasing: bool = False) -> _nt.Array2D[np.complex128]: ... @overload def vander(x: _nt.ToObject_1d, N: int | None = None, increasing: bool = False) -> _nt.Array2D[np.object_]: ... From 0a3cab06bfd20292a9666f8d99b39528a4400d0d Mon Sep 17 00:00:00 2001 From: jorenham Date: Sat, 27 Dec 2025 14:21:43 +0100 Subject: [PATCH 8/9] =?UTF-8?q?=F0=9F=91=BD=EF=B8=8F=20`histogram2d`=20syn?= =?UTF-8?q?c=20with=20upstream?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/lib/_twodim_base_impl.pyi | 72 +++++++++++++++-------- 1 file changed, 49 insertions(+), 23 deletions(-) diff --git a/src/numpy-stubs/lib/_twodim_base_impl.pyi b/src/numpy-stubs/lib/_twodim_base_impl.pyi index 81351397..9bc7049d 100644 --- a/src/numpy-stubs/lib/_twodim_base_impl.pyi +++ b/src/numpy-stubs/lib/_twodim_base_impl.pyi @@ -1,6 +1,6 @@ from _typeshed import Incomplete from collections.abc import Callable, Sequence -from typing import Literal as L, Never, TypeAlias, overload +from typing import Any, Literal as L, Never, TypeAlias, overload from typing_extensions import TypeVar import _numtype as _nt @@ -44,7 +44,6 @@ _MaskFunc: TypeAlias = Callable[ [_nt.Array[np.intp], _T], _nt.Array[_nt.co_complex | np.datetime64 | np.timedelta64 | np.object_] ] -_Sequence01D: TypeAlias = _T | Sequence[_T] _Indices2D: TypeAlias = tuple[_nt.Array1D[np.intp], _nt.Array1D[np.intp]] _Histogram2D: TypeAlias = tuple[_nt.Array2D[np.float64], _nt.Array1D[_ScalarT], _nt.Array1D[_ScalarT]] @@ -204,7 +203,7 @@ def vander(x: _nt.ToObject_1d, N: int | None = None, increasing: bool = False) - def histogram2d( x: _nt._ToArray_1d[_ComplexT], y: _nt._ToArray_1d[_ComplexT | _nt.co_float], - bins: _Sequence01D[int] = 10, + bins: int | Sequence[int] = 10, range: _nt.ToFloat64_2d | None = None, density: bool | None = None, weights: _nt.CoFloat64_1d | None = None, @@ -213,7 +212,7 @@ def histogram2d( def histogram2d( x: _nt._ToArray_1d[_ComplexT | _nt.co_float], y: _nt._ToArray_1d[_ComplexT], - bins: _Sequence01D[int] = 10, + bins: int | Sequence[int] = 10, range: _nt.ToFloat64_2d | None = None, density: bool | None = None, weights: _nt.CoFloat64_1d | None = None, @@ -222,7 +221,7 @@ def histogram2d( def histogram2d( x: _nt._ToArray_1d[_InexactT], y: _nt._ToArray_1d[_InexactT | _nt.co_integer], - bins: _Sequence01D[int] = 10, + bins: int | Sequence[int] = 10, range: _nt.ToFloat64_2d | None = None, density: bool | None = None, weights: _nt.CoFloat64_1d | None = None, @@ -231,16 +230,16 @@ def histogram2d( def histogram2d( x: _nt._ToArray_1d[_InexactT | _nt.co_integer], y: _nt._ToArray_1d[_InexactT], - bins: _Sequence01D[int] = 10, + bins: int | Sequence[int] = 10, range: _nt.ToFloat64_2d | None = None, density: bool | None = None, weights: _nt.CoFloat64_1d | None = None, ) -> _Histogram2D[_InexactT]: ... @overload def histogram2d( - x: _nt.CoFloat64_1d, - y: _nt.CoFloat64_1d, - bins: _Sequence01D[int] = 10, + x: _nt.ToInt_1d | Sequence[float], + y: _nt.ToInt_1d | Sequence[float], + bins: int | Sequence[int] = 10, range: _nt.ToFloat64_2d | None = None, density: bool | None = None, weights: _nt.CoFloat64_1d | None = None, @@ -249,16 +248,16 @@ def histogram2d( def histogram2d( x: Sequence[complex], y: Sequence[complex], - bins: _Sequence01D[int] = 10, + bins: int | Sequence[int] = 10, range: _nt.ToFloat64_2d | None = None, density: bool | None = None, weights: _nt.CoFloat64_1d | None = None, -) -> _Histogram2D[_nt.inexact64]: ... +) -> _Histogram2D[np.complex128 | Any]: ... @overload def histogram2d( x: _nt.CoComplex_1d, y: _nt.CoComplex_1d, - bins: _Sequence01D[_nt._ToArray_1d[_CoComplexT]], + bins: _nt._ToArray_1d[_CoComplexT] | Sequence[_nt._ToArray_1d[_CoComplexT]], range: _nt.ToFloat64_2d | None = None, density: bool | None = None, weights: _nt.CoFloat64_1d | None = None, @@ -271,7 +270,16 @@ def histogram2d( range: _nt.ToFloat64_2d | None = None, density: bool | None = None, weights: _nt.CoFloat64_1d | None = None, -) -> _Histogram2D[_CoComplexT | _InexactT]: ... +) -> _Histogram2D[_InexactT | _CoComplexT]: ... +@overload +def histogram2d( + x: _nt._ToArray_1d[_InexactT], + y: _nt._ToArray_1d[_InexactT], + bins: Sequence[_nt.CoComplex_1d | int], + range: _nt.ToFloat64_2d | None = None, + density: bool | None = None, + weights: _nt.CoFloat64_1d | None = None, +) -> _Histogram2D[_InexactT | Any]: ... @overload def histogram2d( x: _nt.ToInt_1d | Sequence[float], @@ -280,7 +288,16 @@ def histogram2d( range: _nt.ToFloat64_2d | None = None, density: bool | None = None, weights: _nt.CoFloat64_1d | None = None, -) -> _Histogram2D[_CoComplexT | np.float64]: ... +) -> _Histogram2D[np.float64 | _CoComplexT]: ... +@overload +def histogram2d( + x: _nt.ToInt_1d | Sequence[float], + y: _nt.ToInt_1d | Sequence[float], + bins: Sequence[_nt.CoComplex_1d | int], + range: _nt.ToFloat64_2d | None = None, + density: bool | None = None, + weights: _nt.CoFloat64_1d | None = None, +) -> _Histogram2D[np.float64 | Any]: ... @overload def histogram2d( x: Sequence[complex], @@ -289,43 +306,52 @@ def histogram2d( range: _nt.ToFloat64_2d | None = None, density: bool | None = None, weights: _nt.CoFloat64_1d | None = None, -) -> _Histogram2D[_CoComplexT | _nt.inexact64]: ... +) -> _Histogram2D[np.complex128 | _CoComplexT]: ... +@overload +def histogram2d( + x: Sequence[complex], + y: Sequence[complex], + bins: Sequence[_nt.CoComplex_1d | int], + range: _nt.ToFloat64_2d | None = None, + density: bool | None = None, + weights: _nt.CoFloat64_1d | None = None, +) -> _Histogram2D[np.complex128 | Any]: ... @overload def histogram2d( x: _nt.CoComplex_1d, y: _nt.CoComplex_1d, - bins: _nt.Sequence2D[bool], + bins: Sequence[Sequence[int]], range: _nt.ToFloat64_2d | None = None, density: bool | None = None, weights: _nt.CoFloat64_1d | None = None, -) -> _Histogram2D[np.bool]: ... +) -> _Histogram2D[np.int_]: ... @overload def histogram2d( x: _nt.CoComplex_1d, y: _nt.CoComplex_1d, - bins: _nt.Sequence2D[_nt.JustInt], + bins: Sequence[Sequence[float]], range: _nt.ToFloat64_2d | None = None, density: bool | None = None, weights: _nt.CoFloat64_1d | None = None, -) -> _Histogram2D[np.intp]: ... +) -> _Histogram2D[np.float64 | Any]: ... @overload def histogram2d( x: _nt.CoComplex_1d, y: _nt.CoComplex_1d, - bins: _nt.Sequence2D[_nt.JustFloat], + bins: Sequence[Sequence[complex]], range: _nt.ToFloat64_2d | None = None, density: bool | None = None, weights: _nt.CoFloat64_1d | None = None, -) -> _Histogram2D[np.float64]: ... +) -> _Histogram2D[np.complex128 | Any]: ... @overload def histogram2d( x: _nt.CoComplex_1d, y: _nt.CoComplex_1d, - bins: _nt.Sequence2D[_nt.JustComplex], + bins: Sequence[_nt.CoComplex_1d | int] | int, range: _nt.ToFloat64_2d | None = None, density: bool | None = None, weights: _nt.CoFloat64_1d | None = None, -) -> _Histogram2D[np.complex128]: ... +) -> _Histogram2D[Any]: ... # NOTE: we're assuming/demanding here the `mask_func` returns # an ndarray of shape `(n, n)`; otherwise there is the possibility From 0d86d8fae0cd1751d1cf0d6b6c7fdc0fee5c2c30 Mon Sep 17 00:00:00 2001 From: jorenham Date: Sat, 27 Dec 2025 14:22:32 +0100 Subject: [PATCH 9/9] =?UTF-8?q?=E2=9C=85=20update=20`lib.=5Ftwodim=5Fbase?= =?UTF-8?q?=5Fimpl`=20type-tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../@test/static/accept/twodim_base.pyi | 88 ++++++++----------- 1 file changed, 36 insertions(+), 52 deletions(-) diff --git a/src/numpy-stubs/@test/static/accept/twodim_base.pyi b/src/numpy-stubs/@test/static/accept/twodim_base.pyi index b9ec066b..27b492b5 100644 --- a/src/numpy-stubs/@test/static/accept/twodim_base.pyi +++ b/src/numpy-stubs/@test/static/accept/twodim_base.pyi @@ -1,4 +1,4 @@ -from typing import assert_type +from typing import Any, TypeAlias, assert_type from typing_extensions import TypeVar import _numtype as _nt @@ -6,6 +6,8 @@ import numpy as np _ScalarT = TypeVar("_ScalarT", bound=np.generic) +_Histogram2D: TypeAlias = tuple[_nt.Array2D[np.float64], _nt.Array1D[_ScalarT], _nt.Array1D[_ScalarT]] + def func1(ar: _nt.Array[_ScalarT], a: int) -> _nt.Array[_ScalarT]: ... def func2(ar: _nt.Array[np.number], a: str) -> _nt.Array[np.float64]: ... @@ -19,6 +21,8 @@ AR_O: _nt.Array[np.object_] AR_LIKE_b: list[bool] AR_LIKE_c: list[complex] +### + assert_type(np.fliplr(AR_b), _nt.Array[np.bool]) assert_type(np.fliplr(AR_LIKE_b), _nt.Array) @@ -30,14 +34,14 @@ assert_type(np.eye(10, M=20, dtype=np.int64), _nt.Array[np.int64]) assert_type(np.eye(10, k=2, dtype=int), _nt.Array) assert_type(np.diag(AR_b), _nt.Array[np.bool]) -assert_type(np.diag(AR_LIKE_b, k=0), _nt.Array) +assert_type(np.diag(AR_LIKE_b, k=0), _nt.Array2D) -assert_type(np.diagflat(AR_b), _nt.Array[np.bool]) -assert_type(np.diagflat(AR_LIKE_b, k=0), _nt.Array) +assert_type(np.diagflat(AR_b), _nt.Array2D[np.bool]) +assert_type(np.diagflat(AR_LIKE_b, k=0), _nt.Array2D) -assert_type(np.tri(10), _nt.Array[np.float64]) -assert_type(np.tri(10, M=20, dtype=np.int64), _nt.Array[np.int64]) -assert_type(np.tri(10, k=2, dtype=int), _nt.Array) +assert_type(np.tri(10), _nt.Array2D[np.float64]) +assert_type(np.tri(10, M=20, dtype=np.int64), _nt.Array2D[np.int64]) +assert_type(np.tri(10, k=2, dtype=int), _nt.Array2D) assert_type(np.tril(AR_b), _nt.Array[np.bool]) assert_type(np.tril(AR_LIKE_b, k=0), _nt.Array[np.bool]) @@ -45,48 +49,28 @@ assert_type(np.tril(AR_LIKE_b, k=0), _nt.Array[np.bool]) assert_type(np.triu(AR_b), _nt.Array[np.bool]) assert_type(np.triu(AR_LIKE_b, k=0), _nt.Array[np.bool]) -assert_type(np.vander(AR_b), _nt.Array[np.intp]) -assert_type(np.vander(AR_u), _nt.Array[np.signedinteger]) -assert_type(np.vander(AR_i, N=2), _nt.Array[np.signedinteger]) -assert_type(np.vander(AR_f, increasing=True), _nt.Array[np.floating]) -assert_type(np.vander(AR_c), _nt.Array[np.complexfloating]) -assert_type(np.vander(AR_O), _nt.Array[np.object_]) - -assert_type( - np.histogram2d(AR_LIKE_c, AR_LIKE_c), - tuple[_nt.Array[np.float64], _nt.Array[np.complex128 | np.float64], _nt.Array[np.complex128 | np.float64]], -) -assert_type(np.histogram2d(AR_i, AR_b), tuple[_nt.Array[np.float64], _nt.Array[np.float64], _nt.Array[np.float64]]) -assert_type(np.histogram2d(AR_f, AR_i), tuple[_nt.Array[np.float64], _nt.Array[np.float64], _nt.Array[np.float64]]) -assert_type(np.histogram2d(AR_i, AR_f), tuple[_nt.Array[np.float64], _nt.Array[np.float64], _nt.Array[np.float64]]) -assert_type( - np.histogram2d(AR_f, AR_c, weights=AR_LIKE_b), - tuple[_nt.Array[np.float64], _nt.Array[np.complex128], _nt.Array[np.complex128]], -) -assert_type( - np.histogram2d(AR_f, AR_c, bins=8), tuple[_nt.Array[np.float64], _nt.Array[np.complex128], _nt.Array[np.complex128]] -) -assert_type( - np.histogram2d(AR_c, AR_f, bins=(8, 5)), - tuple[_nt.Array[np.float64], _nt.Array[np.complex128], _nt.Array[np.complex128]], -) -assert_type( - np.histogram2d(AR_c, AR_i, bins=AR_u), tuple[_nt.Array[np.float64], _nt.Array[np.uint64], _nt.Array[np.uint64]] -) -assert_type( - np.histogram2d(AR_c, AR_c, bins=(AR_u, AR_u)), - tuple[_nt.Array[np.float64], _nt.Array[np.uint64], _nt.Array[np.uint64]], -) -assert_type( - np.histogram2d(AR_c, AR_c, bins=(AR_b, 8)), - tuple[_nt.Array[np.float64], _nt.Array[np.bool | np.complex128], _nt.Array[np.bool | np.complex128]], -) - -assert_type(np.mask_indices(10, func1), tuple[_nt.Array[np.intp], _nt.Array[np.intp]]) -assert_type(np.mask_indices(8, func2, "0"), tuple[_nt.Array[np.intp], _nt.Array[np.intp]]) - -assert_type(np.tril_indices(10), tuple[_nt.Array[np.intp], _nt.Array[np.intp]]) -assert_type(np.triu_indices(10), tuple[_nt.Array[np.intp], _nt.Array[np.intp]]) - -assert_type(np.tril_indices_from(AR_b), tuple[_nt.Array[np.intp], _nt.Array[np.intp]]) -assert_type(np.triu_indices_from(AR_b), tuple[_nt.Array[np.intp], _nt.Array[np.intp]]) +assert_type(np.vander(AR_b), _nt.Array2D[np.intp]) +assert_type(np.vander(AR_u), _nt.Array2D[np.int_]) +assert_type(np.vander(AR_i, N=2), _nt.Array2D[np.int_]) +assert_type(np.vander(AR_f, increasing=True), _nt.Array2D[np.float64]) +assert_type(np.vander(AR_c), _nt.Array2D[np.complex128]) +assert_type(np.vander(AR_O), _nt.Array2D[np.object_]) + +assert_type(np.histogram2d(AR_LIKE_c, AR_LIKE_c), _Histogram2D[np.complex128 | Any]) +assert_type(np.histogram2d(AR_f, AR_i), _Histogram2D[np.float64]) +assert_type(np.histogram2d(AR_i, AR_f), _Histogram2D[np.float64]) +assert_type(np.histogram2d(AR_f, AR_c, weights=AR_LIKE_b), _Histogram2D[np.complex128]) +assert_type(np.histogram2d(AR_f, AR_c, bins=8), _Histogram2D[np.complex128]) +assert_type(np.histogram2d(AR_c, AR_f, bins=(8, 5)), _Histogram2D[np.complex128]) +assert_type(np.histogram2d(AR_c, AR_i, bins=AR_u), _Histogram2D[np.uint64]) +assert_type(np.histogram2d(AR_c, AR_c, bins=(AR_u, AR_u)), _Histogram2D[np.uint64]) +assert_type(np.histogram2d(AR_c, AR_c, bins=(AR_b, 8)), _Histogram2D[np.bool | np.complex128]) + +assert_type(np.mask_indices(10, func1), tuple[_nt.Array1D[np.intp], _nt.Array1D[np.intp]]) +assert_type(np.mask_indices(8, func2, "0"), tuple[_nt.Array1D[np.intp], _nt.Array1D[np.intp]]) + +assert_type(np.tril_indices(10), tuple[_nt.Array1D[np.intp], _nt.Array1D[np.intp]]) +assert_type(np.triu_indices(10), tuple[_nt.Array1D[np.intp], _nt.Array1D[np.intp]]) + +assert_type(np.tril_indices_from(AR_b), tuple[_nt.Array1D[np.intp], _nt.Array1D[np.intp]]) +assert_type(np.triu_indices_from(AR_b), tuple[_nt.Array1D[np.intp], _nt.Array1D[np.intp]])