diff --git a/scipy-stubs/fft/_helper.pyi b/scipy-stubs/fft/_helper.pyi index 4d119cef..01109f8b 100644 --- a/scipy-stubs/fft/_helper.pyi +++ b/scipy-stubs/fft/_helper.pyi @@ -4,7 +4,7 @@ from typing import Any, Literal, TypeVar, overload import numpy as np import optype as op import optype.numpy as onp -from numpy._typing import _ArrayLike, _NestedSequence +from numpy._typing import _ArrayLike from scipy._typing import AnyShape _SCT = TypeVar("_SCT", bound=np.inexact[Any]) @@ -41,21 +41,21 @@ def rfftfreq(n: int, d: float = 1.0, *, xp: ModuleType, device: object | None = # TODO(jorenham): Array API support (for `x`) # https://github.com/jorenham/scipy-stubs/issues/140 @overload -def fftshift(x: _ArrayLike[_SCT], axes: AnyShape | None = None) -> onp.ArrayND[_SCT]: ... +def fftshift(x: onp.ToIntND | onp.SequenceND[float], axes: AnyShape | None = None) -> onp.ArrayND[np.float64]: ... @overload -def fftshift(x: onp.ToIntND | _NestedSequence[float], axes: AnyShape | None = None) -> onp.ArrayND[np.float64]: ... +def fftshift(x: onp.SequenceND[complex], axes: AnyShape | None = None) -> onp.ArrayND[np.complex128 | np.float64]: ... @overload -def fftshift(x: _NestedSequence[complex], axes: AnyShape | None = None) -> onp.ArrayND[np.complex128 | np.float64]: ... +def fftshift(x: _ArrayLike[_SCT], axes: AnyShape | None = None) -> onp.ArrayND[_SCT]: ... @overload def fftshift(x: onp.ToComplexND, axes: AnyShape | None = None) -> onp.ArrayND[np.inexact[Any]]: ... # TODO(jorenham): Array API support (for `x`) # https://github.com/jorenham/scipy-stubs/issues/140 @overload -def ifftshift(x: _ArrayLike[_SCT], axes: AnyShape | None = None) -> onp.ArrayND[_SCT]: ... +def ifftshift(x: onp.ToIntND | onp.SequenceND[float], axes: AnyShape | None = None) -> onp.ArrayND[np.float64]: ... @overload -def ifftshift(x: onp.ToIntND | _NestedSequence[float], axes: AnyShape | None = None) -> onp.ArrayND[np.float64]: ... +def ifftshift(x: onp.SequenceND[complex], axes: AnyShape | None = None) -> onp.ArrayND[np.complex128 | np.float64]: ... @overload -def ifftshift(x: _NestedSequence[complex], axes: AnyShape | None = None) -> onp.ArrayND[np.complex128 | np.float64]: ... +def ifftshift(x: _ArrayLike[_SCT], axes: AnyShape | None = None) -> onp.ArrayND[_SCT]: ... @overload def ifftshift(x: onp.ToComplexND, axes: AnyShape | None = None) -> onp.ArrayND[np.inexact[Any]]: ... diff --git a/scipy-stubs/signal/_bsplines.pyi b/scipy-stubs/signal/_bsplines.pyi index c4fa6f36..f01f3ca4 100644 --- a/scipy-stubs/signal/_bsplines.pyi +++ b/scipy-stubs/signal/_bsplines.pyi @@ -10,13 +10,11 @@ __all__ = ["cspline1d", "cspline1d_eval", "gauss_spline", "qspline1d", "qspline1 _SCT_fc = TypeVar("_SCT_fc", bound=np.inexact[Any]) @overload -def gauss_spline(x: _ArrayLike[_SCT_fc], n: onp.ToInt) -> onp.ArrayND[_SCT_fc]: ... -@overload -def gauss_spline(x: onp.ToIntND, n: onp.ToInt) -> onp.ArrayND[np.float64]: ... +def gauss_spline(x: onp.ToIntND | onp.SequenceND[float], n: onp.ToInt) -> onp.ArrayND[np.float64]: ... @overload -def gauss_spline(x: onp.ToFloatND, n: onp.ToInt) -> onp.ArrayND[np.floating[Any]]: ... +def gauss_spline(x: onp.SequenceND[complex], n: onp.ToInt) -> onp.ArrayND[np.complex128 | np.float64]: ... @overload -def gauss_spline(x: onp.ToComplexND, n: onp.ToInt) -> onp.ArrayND[np.inexact[Any]]: ... +def gauss_spline(x: _ArrayLike[_SCT_fc], n: onp.ToInt) -> onp.ArrayND[_SCT_fc]: ... # def spline_filter(Iin: _ArrayLike[_SCT_fc], lmbda: onp.ToFloat = 5.0) -> onp.ArrayND[_SCT_fc]: ... diff --git a/scipy-stubs/stats/_mstats_basic.pyi b/scipy-stubs/stats/_mstats_basic.pyi index b7356406..122aacf3 100644 --- a/scipy-stubs/stats/_mstats_basic.pyi +++ b/scipy-stubs/stats/_mstats_basic.pyi @@ -5,6 +5,7 @@ from typing_extensions import Self, TypeVar import numpy as np import optype as op import optype.numpy as onp +import optype.typing as opt from numpy._typing import _ArrayLike from scipy._typing import Alternative, AnyBool, NanPolicy from ._stats_mstats_common import SiegelslopesResult, TheilslopesResult @@ -249,120 +250,162 @@ def kstest( # @overload def trima( - a: _ArrayLike[_SCT_bifc], - limits: tuple[onp.ToComplex, onp.ToComplex] | None = None, + a: onp.SequenceND[bool], + limits: tuple[onp.ToInt, onp.ToInt] | None = None, inclusive: tuple[op.CanBool, op.CanBool] = (True, True), -) -> _MArrayND[_SCT_bifc]: ... +) -> _MArrayND[np.bool_]: ... @overload def trima( - a: onp.ToFloatND, + a: onp.SequenceND[opt.JustInt], + limits: tuple[onp.ToInt, onp.ToInt] | None = None, + inclusive: tuple[op.CanBool, op.CanBool] = (True, True), +) -> _MArrayND[np.int_]: ... +@overload +def trima( + a: onp.SequenceND[float], limits: tuple[onp.ToFloat, onp.ToFloat] | None = None, inclusive: tuple[op.CanBool, op.CanBool] = (True, True), ) -> _MArrayND[np.float64 | np.int_ | np.bool_]: ... @overload def trima( - a: onp.ToComplexND, + a: onp.SequenceND[complex], limits: tuple[onp.ToComplex, onp.ToComplex] | None = None, inclusive: tuple[op.CanBool, op.CanBool] = (True, True), ) -> _MArrayND[np.complex128 | np.float64 | np.int_ | np.bool_]: ... +@overload +def trima( + a: _ArrayLike[_SCT_bifc], + limits: tuple[onp.ToComplex, onp.ToComplex] | None = None, + inclusive: tuple[op.CanBool, op.CanBool] = (True, True), +) -> _MArrayND[_SCT_bifc]: ... # @overload def trimr( - a: _ArrayLike[_SCT_bifc], - limits: tuple[onp.ToComplex, onp.ToComplex] | None = None, + a: onp.SequenceND[opt.JustInt | np.int_], + limits: tuple[onp.ToFloat, onp.ToFloat] | None = None, inclusive: tuple[op.CanBool, op.CanBool] = (True, True), axis: op.CanIndex | None = None, -) -> _MArrayND[_SCT_bifc]: ... +) -> _MArrayND[np.int_]: ... @overload def trimr( - a: onp.ToFloatND, + a: onp.SequenceND[float], limits: tuple[onp.ToFloat, onp.ToFloat] | None = None, inclusive: tuple[op.CanBool, op.CanBool] = (True, True), axis: op.CanIndex | None = None, -) -> _MArrayND[np.float64 | np.int_ | np.bool_]: ... +) -> _MArrayND[np.float64 | np.int_]: ... @overload def trimr( - a: onp.ToComplexND, + a: onp.SequenceND[complex], limits: tuple[onp.ToComplex, onp.ToComplex] | None = None, inclusive: tuple[op.CanBool, op.CanBool] = (True, True), axis: op.CanIndex | None = None, -) -> _MArrayND[np.complex128 | np.float64 | np.int_ | np.bool_]: ... +) -> _MArrayND[np.complex128 | np.float64 | np.int_]: ... +@overload +def trimr( + a: _ArrayLike[_SCT_bifc], + limits: tuple[onp.ToComplex, onp.ToComplex] | None = None, + inclusive: tuple[op.CanBool, op.CanBool] = (True, True), + axis: op.CanIndex | None = None, +) -> _MArrayND[_SCT_bifc]: ... # @overload def trim( - a: _ArrayLike[_SCT_bifc], - limits: tuple[onp.ToComplex, onp.ToComplex] | None = None, + a: onp.SequenceND[opt.JustInt | np.int_], + limits: tuple[onp.ToFloat, onp.ToFloat] | None = None, inclusive: tuple[op.CanBool, op.CanBool] = (True, True), relative: op.CanBool = False, axis: op.CanIndex | None = None, -) -> _MArrayND[_SCT_bifc]: ... +) -> _MArrayND[np.int_]: ... @overload def trim( - a: onp.ToFloatND, + a: onp.SequenceND[float], limits: tuple[onp.ToFloat, onp.ToFloat] | None = None, inclusive: tuple[op.CanBool, op.CanBool] = (True, True), relative: op.CanBool = False, axis: op.CanIndex | None = None, -) -> _MArrayND[np.float64 | np.int_ | np.bool_]: ... +) -> _MArrayND[np.float64 | np.int_]: ... @overload def trim( - a: onp.ToComplexND, + a: onp.SequenceND[complex], limits: tuple[onp.ToComplex, onp.ToComplex] | None = None, inclusive: tuple[op.CanBool, op.CanBool] = (True, True), relative: op.CanBool = False, axis: op.CanIndex | None = None, -) -> _MArrayND[np.complex128 | np.float64 | np.int_ | np.bool_]: ... +) -> _MArrayND[np.complex128 | np.float64 | np.int_]: ... +@overload +def trim( + a: _ArrayLike[_SCT_bifc], + limits: tuple[onp.ToComplex, onp.ToComplex] | None = None, + inclusive: tuple[op.CanBool, op.CanBool] = (True, True), + relative: op.CanBool = False, + axis: op.CanIndex | None = None, +) -> _MArrayND[_SCT_bifc]: ... # @overload def trimboth( - data: _ArrayLike[_SCT_bifc], + data: onp.SequenceND[opt.JustInt | np.int_], proportiontocut: float | np.floating[Any] = 0.2, inclusive: tuple[op.CanBool, op.CanBool] = (True, True), axis: op.CanIndex | None = None, -) -> _MArrayND[_SCT_bifc]: ... +) -> _MArrayND[np.int_]: ... @overload def trimboth( - data: onp.ToFloatND, + data: onp.SequenceND[float], proportiontocut: float | np.floating[Any] = 0.2, inclusive: tuple[op.CanBool, op.CanBool] = (True, True), axis: op.CanIndex | None = None, -) -> _MArrayND[np.float64 | np.int_ | np.bool_]: ... +) -> _MArrayND[np.float64 | np.int_]: ... @overload def trimboth( - data: onp.ToComplexND, + data: onp.SequenceND[complex], proportiontocut: float | np.floating[Any] = 0.2, inclusive: tuple[op.CanBool, op.CanBool] = (True, True), axis: op.CanIndex | None = None, -) -> _MArrayND[np.complex128 | np.float64 | np.int_ | np.bool_]: ... +) -> _MArrayND[np.complex128 | np.float64 | np.int_]: ... +@overload +def trimboth( + data: _ArrayLike[_SCT_bifc], + proportiontocut: float | np.floating[Any] = 0.2, + inclusive: tuple[op.CanBool, op.CanBool] = (True, True), + axis: op.CanIndex | None = None, +) -> _MArrayND[_SCT_bifc]: ... # @overload def trimtail( - data: _ArrayLike[_SCT_bifc], + data: onp.SequenceND[opt.JustInt | np.int_], proportiontocut: float | np.floating[Any] = 0.2, tail: Literal["left", "right"] = "left", inclusive: tuple[op.CanBool, op.CanBool] = (True, True), axis: op.CanIndex | None = None, -) -> _MArrayND[_SCT_bifc]: ... +) -> _MArrayND[np.int_]: ... @overload def trimtail( - data: onp.ToFloatND, + data: onp.SequenceND[float], proportiontocut: float | np.floating[Any] = 0.2, tail: Literal["left", "right"] = "left", inclusive: tuple[op.CanBool, op.CanBool] = (True, True), axis: op.CanIndex | None = None, -) -> _MArrayND[np.float64 | np.int_ | np.bool_]: ... +) -> _MArrayND[np.float64 | np.int_]: ... @overload def trimtail( - data: onp.ToComplexND, + data: onp.SequenceND[complex], proportiontocut: float | np.floating[Any] = 0.2, tail: Literal["left", "right"] = "left", inclusive: tuple[op.CanBool, op.CanBool] = (True, True), axis: op.CanIndex | None = None, -) -> _MArrayND[np.complex128 | np.float64 | np.int_ | np.bool_]: ... +) -> _MArrayND[np.complex128 | np.float64 | np.int_]: ... +@overload +def trimtail( + data: _ArrayLike[_SCT_bifc], + proportiontocut: float | np.floating[Any] = 0.2, + tail: Literal["left", "right"] = "left", + inclusive: tuple[op.CanBool, op.CanBool] = (True, True), + axis: op.CanIndex | None = None, +) -> _MArrayND[_SCT_bifc]: ... # @overload @@ -381,6 +424,8 @@ def trimmed_mean( relative: op.CanBool = True, axis: op.CanIndex | None = None, ) -> _MArrayND0[np.floating[Any] | np.complex128]: ... + +# def trimmed_var( a: onp.ToComplexND, limits: tuple[onp.ToFloat, onp.ToFloat] = (0.1, 0.1), @@ -389,6 +434,8 @@ def trimmed_var( axis: op.CanIndex | None = None, ddof: onp.ToInt = 0, ) -> _MArrayND0[np.float64]: ... + +# def trimmed_std( a: onp.ToComplexND, limits: tuple[onp.ToFloat, onp.ToFloat] = (0.1, 0.1), @@ -397,12 +444,16 @@ def trimmed_std( axis: op.CanIndex | None = None, ddof: onp.ToInt = 0, ) -> _MArrayND0[np.float64]: ... + +# def trimmed_stde( a: onp.ToComplexND, limits: tuple[onp.ToFloat, onp.ToFloat] = (0.1, 0.1), inclusive: tuple[op.CanBool, op.CanBool] = (1, 1), axis: op.CanIndex | None = None, ) -> _MArrayND0[np.float64]: ... + +# @overload def tmean( a: onp.ToFloatND, @@ -417,6 +468,8 @@ def tmean( inclusive: tuple[op.CanBool, op.CanBool] = (True, True), axis: op.CanIndex | None = None, ) -> _MArrayND0[np.inexact[Any]]: ... + +# def tvar( a: _MArrayND, limits: tuple[onp.ToFloat, onp.ToFloat] | None = None, @@ -424,48 +477,68 @@ def tvar( axis: op.CanIndex | None = 0, ddof: onp.ToInt = 1, ) -> _MArrayND0[np.floating[Any]]: ... + +# @overload def tmin( - a: _ArrayLike[_SCT_bifc], - lowerlimit: onp.ToComplex | None = None, + a: onp.SequenceND[opt.JustInt | np.int_], + upperlimit: onp.ToFloat | None = None, axis: op.CanIndex | None = 0, inclusive: AnyBool = True, -) -> _MArrayND0[_SCT_bifc]: ... +) -> _MArrayND0[np.int_]: ... @overload def tmin( - a: onp.ToFloatND, - lowerlimit: onp.ToFloat | None = None, + a: onp.SequenceND[float], + upperlimit: onp.ToFloat | None = None, axis: op.CanIndex | None = 0, inclusive: AnyBool = True, -) -> _MArrayND0[np.float64 | np.int_ | np.bool_]: ... +) -> _MArrayND0[np.float64 | np.int_]: ... @overload def tmin( - a: onp.ToComplexND, - lowerlimit: onp.ToComplex | None = None, + a: onp.SequenceND[complex], + upperlimit: onp.ToComplex | None = None, axis: op.CanIndex | None = 0, inclusive: AnyBool = True, -) -> _MArrayND0[np.complex128 | np.float64 | np.int_ | np.bool_]: ... +) -> _MArrayND0[np.complex128 | np.float64 | np.int_]: ... @overload -def tmax( +def tmin( a: _ArrayLike[_SCT_bifc], upperlimit: onp.ToComplex | None = None, axis: op.CanIndex | None = 0, inclusive: AnyBool = True, ) -> _MArrayND0[_SCT_bifc]: ... + +# @overload def tmax( - a: onp.ToFloatND, + a: onp.SequenceND[opt.JustInt | np.int_], upperlimit: onp.ToFloat | None = None, axis: op.CanIndex | None = 0, inclusive: AnyBool = True, -) -> _MArrayND0[np.float64 | np.int_ | np.bool_]: ... +) -> _MArrayND0[np.int_]: ... @overload def tmax( - a: onp.ToComplexND, + a: onp.SequenceND[float], + upperlimit: onp.ToFloat | None = None, + axis: op.CanIndex | None = 0, + inclusive: AnyBool = True, +) -> _MArrayND0[np.float64 | np.int_]: ... +@overload +def tmax( + a: onp.SequenceND[complex], upperlimit: onp.ToComplex | None = None, axis: op.CanIndex | None = 0, inclusive: AnyBool = True, -) -> _MArrayND0[np.complex128 | np.float64 | np.int_ | np.bool_]: ... +) -> _MArrayND0[np.complex128 | np.float64 | np.int_]: ... +@overload +def tmax( + a: _ArrayLike[_SCT_bifc], + upperlimit: onp.ToComplex | None = None, + axis: op.CanIndex | None = 0, + inclusive: AnyBool = True, +) -> _MArrayND0[_SCT_bifc]: ... + +# def tsem( a: onp.ToComplexND, limits: tuple[onp.ToFloat, onp.ToFloat] | None = None, @@ -477,31 +550,22 @@ def tsem( # @overload def winsorize( - a: _ArrayLike[_SCT_f], - limits: tuple[onp.ToFloat, onp.ToFloat] | None = None, - inclusive: tuple[op.CanBool, op.CanBool] = (True, True), - inplace: AnyBool = False, - axis: op.CanIndex | None = None, - nan_policy: NanPolicy = "propagate", -) -> _MArrayND[_SCT_f]: ... -@overload -def winsorize( - a: onp.ToBoolND, + a: onp.ToIntND, limits: tuple[onp.ToFloat, onp.ToFloat] | None = None, inclusive: tuple[op.CanBool, op.CanBool] = (True, True), inplace: AnyBool = False, axis: op.CanIndex | None = None, nan_policy: NanPolicy = "propagate", -) -> _MArrayND[np.bool_]: ... +) -> _MArrayND[np.int_]: ... @overload def winsorize( - a: onp.ToIntND, + a: _ArrayLike[_SCT_f], limits: tuple[onp.ToFloat, onp.ToFloat] | None = None, inclusive: tuple[op.CanBool, op.CanBool] = (True, True), inplace: AnyBool = False, axis: op.CanIndex | None = None, nan_policy: NanPolicy = "propagate", -) -> _MArrayND[np.bool_ | np.int_]: ... +) -> _MArrayND[_SCT_f]: ... @overload def winsorize( a: onp.ToFloatND, @@ -510,7 +574,7 @@ def winsorize( inplace: AnyBool = False, axis: op.CanIndex | None = None, nan_policy: NanPolicy = "propagate", -) -> _MArrayND[np.bool_ | np.int_ | np.floating[Any]]: ... +) -> _MArrayND[np.floating[Any] | np.int_]: ... @overload def winsorize( a: onp.ToComplexND, @@ -519,7 +583,7 @@ def winsorize( inplace: AnyBool = False, axis: op.CanIndex | None = None, nan_policy: NanPolicy = "propagate", -) -> _MArrayND[np.bool_ | np.int_ | np.floating[Any] | np.complex128]: ... +) -> _MArrayND[np.complex128 | np.floating[Any] | np.int_]: ... # TODO(jorenham): Overloads for complex array-likes def moment( diff --git a/scipy-stubs/stats/_variation.pyi b/scipy-stubs/stats/_variation.pyi index 3238d6a6..9e34bce2 100644 --- a/scipy-stubs/stats/_variation.pyi +++ b/scipy-stubs/stats/_variation.pyi @@ -16,72 +16,72 @@ _SCT_fc = TypeVar("_SCT_fc", bound=np.inexact[Any]) # NOTE: This demonstrates the ridiculous complexity that's required to properly annotate this simple function with "array-likes". # NOTE: Shape-typing hasn't even been included, as that would require even more overloads. -# array-like input with known floating- or complex-floating dtypes +# sequences of `builtins.float`, that implicitly (and inevitably) also cover `builtins.int` and `builtins.bool` @overload def variation( - a: _ArrayLike[_SCT_fc], - axis: None, + a: Sequence[float], + axis: op.CanIndex | None = 0, nan_policy: NanPolicy = "propagate", ddof: onp.ToInt = 0, *, keepdims: Literal[0, False] = False, -) -> _SCT_fc: ... +) -> np.float64: ... @overload def variation( - a: _ArrayLike[_SCT_fc], - axis: op.CanIndex | None = 0, + a: onp.ToIntND | _NestedSequence[float], + axis: None, nan_policy: NanPolicy = "propagate", ddof: onp.ToInt = 0, *, - keepdims: Literal[1, True], -) -> onp.ArrayND[_SCT_fc]: ... + keepdims: Literal[0, False] = False, +) -> np.float64: ... @overload def variation( - a: _ArrayLike[_SCT_fc], + a: onp.ToIntND | _NestedSequence[float], axis: op.CanIndex | None = 0, nan_policy: NanPolicy = "propagate", ddof: onp.ToInt = 0, *, - keepdims: AnyBool = False, -) -> _SCT_fc | onp.ArrayND[_SCT_fc]: ... - -# sequences of `builtins.float`, that implicitly (and inevitably) also cover `builtins.int` and `builtins.bool` + keepdims: Literal[1, True], +) -> onp.ArrayND[np.float64]: ... @overload def variation( - a: Sequence[float], + a: onp.ToIntND | _NestedSequence[float], axis: op.CanIndex | None = 0, nan_policy: NanPolicy = "propagate", ddof: onp.ToInt = 0, *, - keepdims: Literal[0, False] = False, -) -> np.float64: ... + keepdims: AnyBool = False, +) -> np.float64 | onp.ArrayND[np.float64]: ... + +# array-like input with known floating- or complex-floating dtypes @overload def variation( - a: onp.ToIntND | _NestedSequence[float], + a: _ArrayLike[_SCT_fc], axis: None, nan_policy: NanPolicy = "propagate", ddof: onp.ToInt = 0, *, keepdims: Literal[0, False] = False, -) -> np.float64: ... +) -> _SCT_fc: ... @overload def variation( - a: onp.ToIntND | _NestedSequence[float], + a: _ArrayLike[_SCT_fc], axis: op.CanIndex | None = 0, nan_policy: NanPolicy = "propagate", ddof: onp.ToInt = 0, *, keepdims: Literal[1, True], -) -> onp.ArrayND[np.float64]: ... +) -> onp.ArrayND[_SCT_fc]: ... @overload def variation( - a: onp.ToIntND | _NestedSequence[float], + a: _ArrayLike[_SCT_fc], axis: op.CanIndex | None = 0, nan_policy: NanPolicy = "propagate", ddof: onp.ToInt = 0, *, keepdims: AnyBool = False, -) -> np.float64 | onp.ArrayND[np.float64]: ... +) -> _SCT_fc | onp.ArrayND[_SCT_fc]: ... # sequences of `builtin.complex`, which behave as if `float <: complex` and therefore "overlaps" with the `builtins.float` # overloads, hence the `complex128 | float64` returns