From 3499cf332caa4a71a9337285f5c4c1dbe0c7bd7d Mon Sep 17 00:00:00 2001 From: jorenham Date: Tue, 15 Jul 2025 16:58:35 +0200 Subject: [PATCH 1/3] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20`linalg`:=20simplified?= =?UTF-8?q?=20`=5Fmisc`=20dtypes,=20and=20explicit=20`Any`=20dtypes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- scipy-stubs/linalg/_misc.pyi | 31 +++++++++++++------------------ 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/scipy-stubs/linalg/_misc.pyi b/scipy-stubs/linalg/_misc.pyi index 20e56264..fd164eb3 100644 --- a/scipy-stubs/linalg/_misc.pyi +++ b/scipy-stubs/linalg/_misc.pyi @@ -1,4 +1,4 @@ -from typing import Literal, TypeAlias, TypeVar, overload +from typing import Any, Literal, TypeAlias, TypeVar, overload import numpy as np import optype as op @@ -11,8 +11,7 @@ __all__ = ["LinAlgError", "LinAlgWarning", "norm"] _Inf: TypeAlias = float _Order: TypeAlias = Literal["fro", "nuc", 0, 1, -1, 2, -2] | _Inf _Axis: TypeAlias = op.CanIndex | tuple[op.CanIndex, op.CanIndex] - -_SubScalar: TypeAlias = np.complex128 | np.float64 | npc.integer | np.bool_ +_SubScalar: TypeAlias = npc.inexact64 | npc.integer | np.bool_ _ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...]) @@ -23,7 +22,7 @@ class LinAlgWarning(RuntimeWarning): ... # NOTE: the mypy errors are false positives (join vs union) @overload # scalar, axis: None = ... -def norm( # type: ignore[overload-overlap] +def norm( a: complex | _SubScalar, ord: _Order | None = None, axis: None = None, @@ -32,11 +31,7 @@ def norm( # type: ignore[overload-overlap] ) -> np.float64: ... @overload # inexact32, axis: None = ... def norm( - a: np.float32 | np.complex64, - ord: _Order | None = None, - axis: None = None, - keepdims: op.CanBool = False, - check_finite: onp.ToBool = True, + a: npc.inexact32, ord: _Order | None = None, axis: None = None, keepdims: op.CanBool = False, check_finite: onp.ToBool = True ) -> np.float32: ... @overload # longdouble, axis: None = ... def norm( @@ -55,7 +50,7 @@ def norm( check_finite: onp.ToBool = True, ) -> np.float64: ... @overload # float64-coercible array, keepdims: True (positional) -def norm( # type: ignore[overload-overlap] +def norm( a: onp.ArrayND[_SubScalar, _ShapeT], ord: _Order | None, axis: _Axis | None, @@ -63,7 +58,7 @@ def norm( # type: ignore[overload-overlap] check_finite: onp.ToBool = True, ) -> onp.ArrayND[np.float64, _ShapeT]: ... @overload # float64-coercible array, keepdims: True (keyword) -def norm( # type: ignore[overload-overlap] +def norm( a: onp.ArrayND[_SubScalar, _ShapeT], ord: _Order | None = None, axis: _Axis | None = None, @@ -72,7 +67,7 @@ def norm( # type: ignore[overload-overlap] check_finite: onp.ToBool = True, ) -> onp.ArrayND[np.float64, _ShapeT]: ... @overload # float64-coercible array-like, keepdims: True (positional) -def norm( # type: ignore[overload-overlap] +def norm( a: onp.SequenceND[onp.CanArrayND[_SubScalar]] | onp.SequenceND[complex | _SubScalar], ord: _Order | None, axis: _Axis | None, @@ -80,7 +75,7 @@ def norm( # type: ignore[overload-overlap] check_finite: onp.ToBool = True, ) -> onp.ArrayND[np.float64]: ... @overload # float64-coercible array-like, keepdims: True (keyword) -def norm( # type: ignore[overload-overlap] +def norm( a: onp.SequenceND[onp.CanArrayND[_SubScalar]] | onp.SequenceND[complex | _SubScalar], ord: _Order | None = None, axis: _Axis | None = None, @@ -90,7 +85,7 @@ def norm( # type: ignore[overload-overlap] ) -> onp.ArrayND[np.float64]: ... @overload # shaped inexact32 array, keepdims: True (positional) def norm( - a: onp.ArrayND[np.float32 | np.complex64, _ShapeT], + a: onp.ArrayND[npc.inexact32, _ShapeT], ord: _Order | None, axis: _Axis | None, keepdims: onp.ToTrue, @@ -106,7 +101,7 @@ def norm( ) -> onp.ArrayND[np.longdouble, _ShapeT]: ... @overload # shaped inexact32 array, keepdims: True (keyword) def norm( - a: onp.ArrayND[np.float32 | np.complex64, _ShapeT], + a: onp.ArrayND[npc.inexact32, _ShapeT], ord: _Order | None = None, axis: _Axis | None = None, *, @@ -124,7 +119,7 @@ def norm( ) -> onp.ArrayND[np.longdouble, _ShapeT]: ... @overload # scalar array-like, keepdims: True (positional) def norm( - a: onp.SequenceND[onp.CanArrayND[np.float32 | np.complex64]] | onp.SequenceND[np.float32 | np.complex64], + a: onp.SequenceND[onp.CanArrayND[npc.inexact32]] | onp.SequenceND[npc.inexact32], ord: _Order | None, axis: _Axis | None, keepdims: onp.ToTrue, @@ -140,7 +135,7 @@ def norm( ) -> onp.ArrayND[np.longdouble]: ... @overload # scalar array-like, keepdims: True (keyword) def norm( - a: onp.SequenceND[onp.CanArrayND[np.float32 | np.complex64]] | onp.SequenceND[np.float32 | np.complex64], + a: onp.SequenceND[onp.CanArrayND[npc.inexact32]] | onp.SequenceND[npc.inexact32], ord: _Order | None = None, axis: _Axis | None = None, *, @@ -187,4 +182,4 @@ def norm( ) -> npc.floating | onp.ArrayND[npc.floating]: ... # -def _datacopied(arr: onp.ArrayND, original: onp.CanArrayND) -> bool: ... # undocumented +def _datacopied(arr: onp.ArrayND[Any], original: onp.CanArrayND[Any]) -> bool: ... # undocumented From 436ad9782708106f4a58f7517b59b719ef73a70e Mon Sep 17 00:00:00 2001 From: jorenham Date: Tue, 15 Jul 2025 17:02:56 +0200 Subject: [PATCH 2/3] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20`fft`:=20shuffle=20(di?= =?UTF-8?q?sjoint)=20`[i]fft*`=20overloads=20to=20avoid=20triggering=20a?= =?UTF-8?q?=20pyright=20bug?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- scipy-stubs/fft/_basic.pyi | 194 +++++++++++++++++++------------------ 1 file changed, 99 insertions(+), 95 deletions(-) diff --git a/scipy-stubs/fft/_basic.pyi b/scipy-stubs/fft/_basic.pyi index 0ba26f7c..434daf8c 100644 --- a/scipy-stubs/fft/_basic.pyi +++ b/scipy-stubs/fft/_basic.pyi @@ -13,24 +13,28 @@ _ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...]) _Norm: TypeAlias = Literal["backward", "ortho", "forward"] _Unused: TypeAlias = Never # not used by scipy +_CoInteger: TypeAlias = npc.integer | np.bool_ + _AsFloat32: TypeAlias = onp.CanArrayND[npc.floating32, _ShapeT] -_AsFloat64: TypeAlias = onp.CanArrayND[np.bool_ | npc.integer | npc.floating64, _ShapeT] +_AsFloat64: TypeAlias = onp.CanArrayND[npc.floating64 | _CoInteger, _ShapeT] _AsFloat80: TypeAlias = onp.CanArrayND[np.longdouble, _ShapeT] -_AsComplex64: TypeAlias = onp.CanArrayND[np.float32 | np.complex64, _ShapeT] -_AsComplex128: TypeAlias = onp.CanArrayND[np.bool_ | npc.integer | np.float64 | np.complex128, _ShapeT] +_AsComplex64: TypeAlias = onp.CanArrayND[npc.inexact32, _ShapeT] +_AsComplex128: TypeAlias = onp.CanArrayND[npc.inexact64 | _CoInteger, _ShapeT] _AsComplex160: TypeAlias = onp.CanArrayND[np.longdouble | np.clongdouble, _ShapeT] -_ToFloat64_ND: TypeAlias = onp.ToArrayND[float, np.bool_ | npc.integer | np.float64] -_ToComplex128_ND: TypeAlias = onp.ToArrayND[complex, np.bool_ | npc.integer | npc.inexact64] +_ToFloat64_ND: TypeAlias = onp.ToArrayND[float, npc.floating64 | _CoInteger] +_ToComplex128_ND: TypeAlias = onp.ToArrayND[complex, npc.inexact64 | _CoInteger] + +# NOTE: The order of overloads is carefully chosen to avoid a pyright bug. ### # 1-D # keep in sync with `ifft` @overload -def fft( # type: ignore[overload-overlap] - x: _AsComplex64[_ShapeT], +def fft( + x: _AsComplex128[_ShapeT], n: int | None = None, axis: int = -1, norm: _Norm | None = None, @@ -38,10 +42,10 @@ def fft( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex64]: ... +) -> onp.Array[_ShapeT, np.complex128]: ... @overload def fft( - x: _AsComplex128[_ShapeT], + x: _AsComplex64[_ShapeT], n: int | None = None, axis: int = -1, norm: _Norm | None = None, @@ -49,7 +53,7 @@ def fft( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex128]: ... +) -> onp.Array[_ShapeT, np.complex64]: ... @overload def fft( x: _AsComplex160[_ShapeT], @@ -97,8 +101,8 @@ def fft( # keep in sync with `fft` @overload -def ifft( # type: ignore[overload-overlap] - x: _AsComplex64[_ShapeT], +def ifft( + x: _AsComplex128[_ShapeT], n: int | None = None, axis: int = -1, norm: _Norm | None = None, @@ -106,10 +110,10 @@ def ifft( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex64]: ... +) -> onp.Array[_ShapeT, np.complex128]: ... @overload def ifft( - x: _AsComplex128[_ShapeT], + x: _AsComplex64[_ShapeT], n: int | None = None, axis: int = -1, norm: _Norm | None = None, @@ -117,7 +121,7 @@ def ifft( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex128]: ... +) -> onp.Array[_ShapeT, np.complex64]: ... @overload def ifft( x: _AsComplex160[_ShapeT], @@ -165,8 +169,8 @@ def ifft( # keep in sync with `ihfft` @overload -def rfft( # type: ignore[overload-overlap] - x: _AsFloat32[_ShapeT], +def rfft( + x: _AsFloat64[_ShapeT], n: int | None = None, axis: int = -1, norm: _Norm | None = None, @@ -174,10 +178,10 @@ def rfft( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex64]: ... +) -> onp.Array[_ShapeT, np.complex128]: ... @overload def rfft( - x: _AsFloat64[_ShapeT], + x: _AsFloat32[_ShapeT], n: int | None = None, axis: int = -1, norm: _Norm | None = None, @@ -185,7 +189,7 @@ def rfft( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex128]: ... +) -> onp.Array[_ShapeT, np.complex64]: ... @overload def rfft( x: _AsFloat80[_ShapeT], @@ -233,8 +237,8 @@ def rfft( # keep in sync with `hfft` @overload -def irfft( # type: ignore[overload-overlap] - x: _AsComplex64[_ShapeT], +def irfft( + x: _AsComplex128[_ShapeT], n: int | None = None, axis: int = -1, norm: _Norm | None = None, @@ -242,10 +246,10 @@ def irfft( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.float32]: ... +) -> onp.Array[_ShapeT, np.float64]: ... @overload def irfft( - x: _AsComplex128[_ShapeT], + x: _AsComplex64[_ShapeT], n: int | None = None, axis: int = -1, norm: _Norm | None = None, @@ -253,7 +257,7 @@ def irfft( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.float64]: ... +) -> onp.Array[_ShapeT, np.float32]: ... @overload def irfft( x: _AsComplex160[_ShapeT], @@ -301,8 +305,8 @@ def irfft( # keep in sync with `irfft` @overload -def hfft( # type: ignore[overload-overlap] - x: _AsComplex64[_ShapeT], +def hfft( + x: _AsComplex128[_ShapeT], n: int | None = None, axis: int = -1, norm: _Norm | None = None, @@ -310,10 +314,10 @@ def hfft( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.float32]: ... +) -> onp.Array[_ShapeT, np.float64]: ... @overload def hfft( - x: _AsComplex128[_ShapeT], + x: _AsComplex64[_ShapeT], n: int | None = None, axis: int = -1, norm: _Norm | None = None, @@ -321,7 +325,7 @@ def hfft( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.float64]: ... +) -> onp.Array[_ShapeT, np.float32]: ... @overload def hfft( x: _AsComplex160[_ShapeT], @@ -370,8 +374,8 @@ def hfft( # # keep in sync with `rfft` @overload -def ihfft( # type: ignore[overload-overlap] - x: _AsFloat32[_ShapeT], +def ihfft( + x: _AsFloat64[_ShapeT], n: int | None = None, axis: int = -1, norm: _Norm | None = None, @@ -379,10 +383,10 @@ def ihfft( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex64]: ... +) -> onp.Array[_ShapeT, np.complex128]: ... @overload def ihfft( - x: _AsFloat64[_ShapeT], + x: _AsFloat32[_ShapeT], n: int | None = None, axis: int = -1, norm: _Norm | None = None, @@ -390,7 +394,7 @@ def ihfft( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex128]: ... +) -> onp.Array[_ShapeT, np.complex64]: ... @overload def ihfft( x: _AsFloat80[_ShapeT], @@ -441,8 +445,8 @@ def ihfft( # keep in sync with `ifft2` @overload -def fft2( # type: ignore[overload-overlap] - x: _AsComplex64[_ShapeT], +def fft2( + x: _AsComplex128[_ShapeT], s: onp.ToJustInt1D | None = None, axes: AnyShape = (-2, -1), norm: _Norm | None = None, @@ -450,10 +454,10 @@ def fft2( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex64]: ... +) -> onp.Array[_ShapeT, np.complex128]: ... @overload def fft2( - x: _AsComplex128[_ShapeT], + x: _AsComplex64[_ShapeT], s: onp.ToJustInt1D | None = None, axes: AnyShape = (-2, -1), norm: _Norm | None = None, @@ -461,7 +465,7 @@ def fft2( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex128]: ... +) -> onp.Array[_ShapeT, np.complex64]: ... @overload def fft2( x: _AsComplex160[_ShapeT], @@ -509,8 +513,8 @@ def fft2( # keep in sync with `fft2` @overload -def ifft2( # type: ignore[overload-overlap] - x: _AsComplex64[_ShapeT], +def ifft2( + x: _AsComplex128[_ShapeT], s: onp.ToJustInt1D | None = None, axes: AnyShape = (-2, -1), norm: _Norm | None = None, @@ -518,10 +522,10 @@ def ifft2( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex64]: ... +) -> onp.Array[_ShapeT, np.complex128]: ... @overload def ifft2( - x: _AsComplex128[_ShapeT], + x: _AsComplex64[_ShapeT], s: onp.ToJustInt1D | None = None, axes: AnyShape = (-2, -1), norm: _Norm | None = None, @@ -529,7 +533,7 @@ def ifft2( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex128]: ... +) -> onp.Array[_ShapeT, np.complex64]: ... @overload def ifft2( x: _AsComplex160[_ShapeT], @@ -577,8 +581,8 @@ def ifft2( # keep in sync with `ihfft2` @overload -def rfft2( # type: ignore[overload-overlap] - x: _AsFloat32[_ShapeT], +def rfft2( + x: _AsFloat64[_ShapeT], s: onp.ToJustInt1D | None = None, axes: AnyShape = (-2, -1), norm: _Norm | None = None, @@ -586,10 +590,10 @@ def rfft2( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex64]: ... +) -> onp.Array[_ShapeT, np.complex128]: ... @overload def rfft2( - x: _AsFloat64[_ShapeT], + x: _AsFloat32[_ShapeT], s: onp.ToJustInt1D | None = None, axes: AnyShape = (-2, -1), norm: _Norm | None = None, @@ -597,7 +601,7 @@ def rfft2( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex128]: ... +) -> onp.Array[_ShapeT, np.complex64]: ... @overload def rfft2( x: _AsFloat80[_ShapeT], @@ -645,8 +649,8 @@ def rfft2( # keep in sync with `hfft2` @overload -def irfft2( # type: ignore[overload-overlap] - x: _AsComplex64[_ShapeT], +def irfft2( + x: _AsComplex128[_ShapeT], s: onp.ToJustInt1D | None = None, axes: AnyShape = (-2, -1), norm: _Norm | None = None, @@ -654,10 +658,10 @@ def irfft2( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.float32]: ... +) -> onp.Array[_ShapeT, np.float64]: ... @overload def irfft2( - x: _AsComplex128[_ShapeT], + x: _AsComplex64[_ShapeT], s: onp.ToJustInt1D | None = None, axes: AnyShape = (-2, -1), norm: _Norm | None = None, @@ -665,7 +669,7 @@ def irfft2( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.float64]: ... +) -> onp.Array[_ShapeT, np.float32]: ... @overload def irfft2( x: _AsComplex160[_ShapeT], @@ -713,8 +717,8 @@ def irfft2( # keep in sync with `irfft2` @overload -def hfft2( # type: ignore[overload-overlap] - x: _AsComplex64[_ShapeT], +def hfft2( + x: _AsComplex128[_ShapeT], s: onp.ToJustInt1D | None = None, axes: AnyShape = (-2, -1), norm: _Norm | None = None, @@ -722,10 +726,10 @@ def hfft2( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.float32]: ... +) -> onp.Array[_ShapeT, np.float64]: ... @overload def hfft2( - x: _AsComplex128[_ShapeT], + x: _AsComplex64[_ShapeT], s: onp.ToJustInt1D | None = None, axes: AnyShape = (-2, -1), norm: _Norm | None = None, @@ -733,7 +737,7 @@ def hfft2( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.float64]: ... +) -> onp.Array[_ShapeT, np.float32]: ... @overload def hfft2( x: _AsComplex160[_ShapeT], @@ -781,8 +785,8 @@ def hfft2( # keep in sync with `rfft2` @overload -def ihfft2( # type: ignore[overload-overlap] - x: _AsFloat32[_ShapeT], +def ihfft2( + x: _AsFloat64[_ShapeT], s: onp.ToJustInt1D | None = None, axes: AnyShape = (-2, -1), norm: _Norm | None = None, @@ -790,10 +794,10 @@ def ihfft2( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex64]: ... +) -> onp.Array[_ShapeT, np.complex128]: ... @overload def ihfft2( - x: _AsFloat64[_ShapeT], + x: _AsFloat32[_ShapeT], s: onp.ToJustInt1D | None = None, axes: AnyShape = (-2, -1), norm: _Norm | None = None, @@ -801,7 +805,7 @@ def ihfft2( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex128]: ... +) -> onp.Array[_ShapeT, np.complex64]: ... @overload def ihfft2( x: _AsFloat80[_ShapeT], @@ -852,8 +856,8 @@ def ihfft2( # keep in sync with `ifftn` @overload -def fftn( # type: ignore[overload-overlap] - x: _AsComplex64[_ShapeT], +def fftn( + x: _AsComplex128[_ShapeT], s: onp.ToJustIntND | None = None, axes: AnyShape | None = None, norm: _Norm | None = None, @@ -861,10 +865,10 @@ def fftn( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex64]: ... +) -> onp.Array[_ShapeT, np.complex128]: ... @overload def fftn( - x: _AsComplex128[_ShapeT], + x: _AsComplex64[_ShapeT], s: onp.ToJustIntND | None = None, axes: AnyShape | None = None, norm: _Norm | None = None, @@ -872,7 +876,7 @@ def fftn( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex128]: ... +) -> onp.Array[_ShapeT, np.complex64]: ... @overload def fftn( x: _AsComplex160[_ShapeT], @@ -909,8 +913,8 @@ def fftn( # keep in sync with `fftn` @overload -def ifftn( # type: ignore[overload-overlap] - x: _AsComplex64[_ShapeT], +def ifftn( + x: _AsComplex128[_ShapeT], s: onp.ToJustIntND | None = None, axes: AnyShape | None = None, norm: _Norm | None = None, @@ -918,10 +922,10 @@ def ifftn( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex64]: ... +) -> onp.Array[_ShapeT, np.complex128]: ... @overload def ifftn( - x: _AsComplex128[_ShapeT], + x: _AsComplex64[_ShapeT], s: onp.ToJustIntND | None = None, axes: AnyShape | None = None, norm: _Norm | None = None, @@ -929,7 +933,7 @@ def ifftn( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex128]: ... +) -> onp.Array[_ShapeT, np.complex64]: ... @overload def ifftn( x: _AsComplex160[_ShapeT], @@ -966,8 +970,8 @@ def ifftn( # keep in sync with `ihfftn` @overload -def rfftn( # type: ignore[overload-overlap] - x: _AsFloat32[_ShapeT], +def rfftn( + x: _AsFloat64[_ShapeT], s: onp.ToJustIntND | None = None, axes: AnyShape | None = None, norm: _Norm | None = None, @@ -975,10 +979,10 @@ def rfftn( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex64]: ... +) -> onp.Array[_ShapeT, np.complex128]: ... @overload def rfftn( - x: _AsFloat64[_ShapeT], + x: _AsFloat32[_ShapeT], s: onp.ToJustIntND | None = None, axes: AnyShape | None = None, norm: _Norm | None = None, @@ -986,7 +990,7 @@ def rfftn( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex128]: ... +) -> onp.Array[_ShapeT, np.complex64]: ... @overload def rfftn( x: _AsFloat80[_ShapeT], @@ -1023,8 +1027,8 @@ def rfftn( # keep in sync with `hfftn` @overload -def irfftn( # type: ignore[overload-overlap] - x: _AsComplex64[_ShapeT], +def irfftn( + x: _AsComplex128[_ShapeT], s: onp.ToJustIntND | None = None, axes: AnyShape | None = None, norm: _Norm | None = None, @@ -1032,10 +1036,10 @@ def irfftn( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.float32]: ... +) -> onp.Array[_ShapeT, np.float64]: ... @overload def irfftn( - x: _AsComplex128[_ShapeT], + x: _AsComplex64[_ShapeT], s: onp.ToJustIntND | None = None, axes: AnyShape | None = None, norm: _Norm | None = None, @@ -1043,7 +1047,7 @@ def irfftn( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.float64]: ... +) -> onp.Array[_ShapeT, np.float32]: ... @overload def irfftn( x: _AsComplex160[_ShapeT], @@ -1080,8 +1084,8 @@ def irfftn( # keep in sync with `irfftn` @overload -def hfftn( # type: ignore[overload-overlap] - x: _AsComplex64[_ShapeT], +def hfftn( + x: _AsComplex128[_ShapeT], s: onp.ToJustIntND | None = None, axes: AnyShape | None = None, norm: _Norm | None = None, @@ -1089,10 +1093,10 @@ def hfftn( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.float32]: ... +) -> onp.Array[_ShapeT, np.float64]: ... @overload def hfftn( - x: _AsComplex128[_ShapeT], + x: _AsComplex64[_ShapeT], s: onp.ToJustIntND | None = None, axes: AnyShape | None = None, norm: _Norm | None = None, @@ -1100,7 +1104,7 @@ def hfftn( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.float64]: ... +) -> onp.Array[_ShapeT, np.float32]: ... @overload def hfftn( x: _AsComplex160[_ShapeT], @@ -1137,8 +1141,8 @@ def hfftn( # keep in sync with `rfftn` @overload -def ihfftn( # type: ignore[overload-overlap] - x: _AsFloat32[_ShapeT], +def ihfftn( + x: _AsFloat64[_ShapeT], s: onp.ToJustIntND | None = None, axes: AnyShape | None = None, norm: _Norm | None = None, @@ -1146,10 +1150,10 @@ def ihfftn( # type: ignore[overload-overlap] workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex64]: ... +) -> onp.Array[_ShapeT, np.complex128]: ... @overload def ihfftn( - x: _AsFloat64[_ShapeT], + x: _AsFloat32[_ShapeT], s: onp.ToJustIntND | None = None, axes: AnyShape | None = None, norm: _Norm | None = None, @@ -1157,7 +1161,7 @@ def ihfftn( workers: int | None = None, *, plan: _Unused | None = None, -) -> onp.Array[_ShapeT, np.complex128]: ... +) -> onp.Array[_ShapeT, np.complex64]: ... @overload def ihfftn( x: _AsFloat80[_ShapeT], From efd159296514f633ee2c046b82fce6f80be063e3 Mon Sep 17 00:00:00 2001 From: jorenham Date: Tue, 15 Jul 2025 17:04:13 +0200 Subject: [PATCH 3/3] =?UTF-8?q?=F0=9F=8E=A8=20`cluster.vq`:=20simplify=20a?= =?UTF-8?q?nd=20shuffle=20overloads=20to=20avoid=20triggering=20a=20pyrigh?= =?UTF-8?q?t=20bug?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- scipy-stubs/cluster/vq.pyi | 61 +++++++++++++++++++++----------------- 1 file changed, 33 insertions(+), 28 deletions(-) diff --git a/scipy-stubs/cluster/vq.pyi b/scipy-stubs/cluster/vq.pyi index 55d6ec3c..315357a4 100644 --- a/scipy-stubs/cluster/vq.pyi +++ b/scipy-stubs/cluster/vq.pyi @@ -1,4 +1,4 @@ -from collections.abc import Callable +from collections.abc import Callable, Sequence from types import ModuleType from typing import Final, Literal, TypeAlias, overload from typing_extensions import TypeVar @@ -9,12 +9,17 @@ import optype.numpy.compat as npc __all__ = ["kmeans", "kmeans2", "vq", "whiten"] +_InexactT = TypeVar("_InexactT", bound=npc.inexact) + _InitMethod: TypeAlias = Literal["random", "points", "++", "matrix"] _MissingMethod: TypeAlias = Literal["warn", "raise"] -_InexactT = TypeVar("_InexactT", bound=npc.inexact) +_ToFloat32_2D: TypeAlias = onp.ToArray2D[int, np.float32 | np.float16 | npc.integer16 | npc.integer8] +_AsFloat64_2D: TypeAlias = onp.ToArray2D[float, npc.floating64 | npc.integer] +_PyFloatMax2D: TypeAlias = Sequence[float] | Sequence[Sequence[float]] ### +# NOTE: DO NOT RE-ORDER THE OVERLOADS WITH a `# type: ignore`, otherwise it'll trigger a pernicious bug in pyright (1.1.403). class ClusterError(Exception): ... @@ -25,14 +30,14 @@ def whiten(obs: onp.ArrayND[np.bool_ | npc.integer], check_finite: bool | None = def whiten(obs: onp.ArrayND[_InexactT], check_finite: bool | None = None) -> onp.Array2D[_InexactT]: ... # +@overload # float32 +def vq( # type: ignore[overload-overlap] + obs: onp.CanArrayND[np.float32], code_book: _ToFloat32_2D, check_finite: bool = True +) -> tuple[onp.Array1D[np.int32], onp.Array1D[np.float32]]: ... @overload # float64 def vq( - obs: onp.ToJustFloat64_2D, code_book: onp.ToInt2D | onp.ToFloat64_2D, check_finite: bool = True + obs: onp.ToJustFloat64_2D, code_book: _AsFloat64_2D, check_finite: bool = True ) -> tuple[onp.Array1D[np.int32], onp.Array1D[np.float64]]: ... -@overload # float32 -def vq( - obs: onp.CanArrayND[np.float32], code_book: onp.ToInt2D | onp.CanArrayND[np.float16 | np.float32], check_finite: bool = True -) -> tuple[onp.Array1D[np.int32], onp.Array1D[np.float32]]: ... @overload # floating def vq( obs: onp.ToJustFloat2D, code_book: onp.ToFloat2D, check_finite: bool = True @@ -49,32 +54,32 @@ def py_vq( ) -> tuple[onp.Array1D[np.intp], onp.Array1D[npc.floating]]: ... # -@overload # float64 -def kmeans( - obs: onp.ToJustFloat64_2D, - k_or_guess: onp.ToJustInt | onp.ToFloat64_ND, +@overload # float32 +def kmeans( # type: ignore[overload-overlap] + obs: onp.CanArrayND[np.float32], + k_or_guess: int | _ToFloat32_2D, iter: int = 20, thresh: float = 1e-5, check_finite: bool = True, *, seed: onp.random.ToRNG | None = None, rng: onp.random.ToRNG | None = None, -) -> tuple[onp.Array2D[np.float64], np.float64]: ... -@overload # float32 +) -> tuple[onp.Array2D[np.float32], np.float32]: ... +@overload # float64 def kmeans( - obs: onp.CanArrayND[np.float32], - k_or_guess: onp.ToJustInt | onp.ToFloatND, + obs: onp.ToJustFloat64_2D, + k_or_guess: int | _AsFloat64_2D, iter: int = 20, thresh: float = 1e-5, check_finite: bool = True, *, seed: onp.random.ToRNG | None = None, rng: onp.random.ToRNG | None = None, -) -> tuple[onp.Array2D[np.float32], np.float32]: ... +) -> tuple[onp.Array2D[np.float64], np.float64]: ... @overload # floating def kmeans( obs: onp.ToJustFloat2D, - k_or_guess: onp.ToJustInt | onp.ToFloatND, + k_or_guess: int | onp.ToFloat2D, iter: int = 20, thresh: float = 1e-5, check_finite: bool = True, @@ -104,10 +109,10 @@ def _missing_raise() -> None: ... # undocumented _valid_miss_meth: Final[dict[str, Callable[[], None]]] = ... # undocumented # -@overload # float64 -def kmeans2( - data: onp.ToJustFloat64_1D | onp.ToJustFloat64_2D, - k: onp.ToJustInt | onp.ToFloatND, +@overload # float32 +def kmeans2( # type: ignore[overload-overlap] + data: onp.CanArrayND[np.float32], + k: int | _ToFloat32_2D, iter: int = 10, thresh: float = 1e-5, minit: _InitMethod = "random", @@ -116,11 +121,11 @@ def kmeans2( *, seed: onp.random.ToRNG | None = None, rng: onp.random.ToRNG | None = None, -) -> tuple[onp.Array2D[np.float64], onp.Array1D[np.int32]]: ... -@overload # float32 +) -> tuple[onp.Array2D[np.float32], onp.Array1D[np.int32]]: ... +@overload # float64 def kmeans2( - data: onp.CanArrayND[np.float32], - k: onp.ToJustInt | onp.ToFloatND, + data: onp.CanArrayND[np.float64] | _PyFloatMax2D, + k: int | _AsFloat64_2D, iter: int = 10, thresh: float = 1e-5, minit: _InitMethod = "random", @@ -129,11 +134,11 @@ def kmeans2( *, seed: onp.random.ToRNG | None = None, rng: onp.random.ToRNG | None = None, -) -> tuple[onp.Array2D[np.float32], onp.Array1D[np.int32]]: ... +) -> tuple[onp.Array2D[np.float64], onp.Array1D[np.int32]]: ... @overload # floating def kmeans2( - data: onp.ToJustFloat1D | onp.ToJustFloat2D, - k: onp.ToJustInt | onp.ToFloatND, + data: onp.CanArrayND[npc.floating] | _PyFloatMax2D, + k: int | onp.ToFloat2D, iter: int = 10, thresh: float = 1e-5, minit: _InitMethod = "random",