diff --git a/scipy-stubs/signal/windows/_windows.pyi b/scipy-stubs/signal/windows/_windows.pyi index b751baa0..7cd32db5 100644 --- a/scipy-stubs/signal/windows/_windows.pyi +++ b/scipy-stubs/signal/windows/_windows.pyi @@ -1,4 +1,11 @@ -from scipy._typing import Untyped +from collections.abc import Sequence +from typing import Any, Literal, TypeAlias, overload +from typing_extensions import Unpack + +import numpy as np +import numpy.typing as npt +import optype as op +from scipy._typing import AnyInt, AnyReal __all__ = [ "barthann", @@ -29,36 +36,180 @@ __all__ = [ "tukey", ] -def general_cosine(M: Untyped, a: Untyped, sym: bool = True) -> Untyped: ... -def boxcar(M: Untyped, sym: bool = True) -> Untyped: ... -def triang(M: Untyped, sym: bool = True) -> Untyped: ... -def parzen(M: Untyped, sym: bool = True) -> Untyped: ... -def bohman(M: Untyped, sym: bool = True) -> Untyped: ... -def blackman(M: Untyped, sym: bool = True) -> Untyped: ... -def nuttall(M: Untyped, sym: bool = True) -> Untyped: ... -def blackmanharris(M: Untyped, sym: bool = True) -> Untyped: ... -def flattop(M: Untyped, sym: bool = True) -> Untyped: ... -def bartlett(M: Untyped, sym: bool = True) -> Untyped: ... -def hann(M: Untyped, sym: bool = True) -> Untyped: ... -def tukey(M: Untyped, alpha: float = 0.5, sym: bool = True) -> Untyped: ... -def barthann(M: Untyped, sym: bool = True) -> Untyped: ... -def general_hamming(M: Untyped, alpha: Untyped, sym: bool = True) -> Untyped: ... -def hamming(M: Untyped, sym: bool = True) -> Untyped: ... -def kaiser(M: Untyped, beta: Untyped, sym: bool = True) -> Untyped: ... -def kaiser_bessel_derived(M: Untyped, beta: Untyped, *, sym: bool = True) -> Untyped: ... -def gaussian(M: Untyped, std: Untyped, sym: bool = True) -> Untyped: ... -def general_gaussian(M: Untyped, p: Untyped, sig: Untyped, sym: bool = True) -> Untyped: ... -def chebwin(M: Untyped, at: Untyped, sym: bool = True) -> Untyped: ... -def cosine(M: Untyped, sym: bool = True) -> Untyped: ... -def exponential(M: Untyped, center: Untyped | None = None, tau: float = 1.0, sym: bool = True) -> Untyped: ... -def taylor(M: Untyped, nbar: int = 4, sll: int = 30, norm: bool = True, sym: bool = True) -> Untyped: ... +_Array_f8_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.float64]] +_Array_f8_2d: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.float64]] +_Weights: TypeAlias = Sequence[AnyReal] | npt.NDArray[np.floating[Any]] | npt.NDArray[np.integer[Any]] | npt.NDArray[np.bool_] + +_Norm: TypeAlias = Literal[2, "approximate", "subsample"] +_WindowLength: TypeAlias = int | np.int16 | np.int32 | np.int64 +_Window: TypeAlias = Literal[ + "barthann", + "brthan", + "bth", + "bartlett", + "bart", + "brt", + "blackman", + "black", + "blk", + "blackmanharris", + "blackharr", + "bkh", + "bohman", + "bman", + "bmn", + "boxcar", + "box", + "ones", + "rect", + "rectangular", + "cosine", + "halfcosine", + "exponential", + "poisson", + "flattop", + "flat", + "flt", + "hamming", + "hamm", + "ham", + "hann", + "han", + "lanczos", + "sinc", + "nuttall", + "nutl", + "nut", + "parzen", + "parz", + "par", + "taylor", + "taylorwin", + "triangle", + "triang", + "tri", + "tukey", + "tuk", +] +_WindowNeedsParams: TypeAlias = Literal[ + "chebwin", + "cheb", + "dpss", + "gaussian", + "gauss", + "gss", + "general cosine", + "general_cosine", + "general gaussian", + "general_gaussian", + "general gauss", + "general_gauss", + "ggs", + "general hamming", + "general_hamming", + "kaiser", + "ksr", + "kaiser bessel derived", + "kbd", +] + +def general_cosine(M: _WindowLength, a: _Weights, sym: op.CanBool = True) -> _Array_f8_1d: ... +def boxcar(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ... +def triang(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ... +def parzen(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ... +def bohman(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ... +def blackman(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ... +def nuttall(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ... +def blackmanharris(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ... +def flattop(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ... +def bartlett(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ... +def hann(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ... +def tukey(M: _WindowLength, alpha: AnyReal = 0.5, sym: op.CanBool = True) -> _Array_f8_1d: ... +def barthann(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ... +def general_hamming(M: _WindowLength, alpha: AnyReal, sym: op.CanBool = True) -> _Array_f8_1d: ... +def hamming(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ... +def kaiser(M: _WindowLength, beta: AnyReal, sym: op.CanBool = True) -> _Array_f8_1d: ... +def kaiser_bessel_derived(M: _WindowLength, beta: AnyReal, *, sym: op.CanBool = True) -> _Array_f8_1d: ... +def gaussian(M: _WindowLength, std: AnyReal, sym: op.CanBool = True) -> _Array_f8_1d: ... +def general_gaussian(M: _WindowLength, p: AnyReal, sig: AnyReal, sym: op.CanBool = True) -> _Array_f8_1d: ... +def chebwin(M: _WindowLength, at: AnyReal, sym: op.CanBool = True) -> _Array_f8_1d: ... +def cosine(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ... +def exponential(M: _WindowLength, center: AnyReal | None = None, tau: AnyReal = 1.0, sym: op.CanBool = True) -> _Array_f8_1d: ... +def taylor( + M: _WindowLength, nbar: AnyInt = 4, sll: AnyInt = 30, norm: op.CanBool = True, sym: op.CanBool = True +) -> _Array_f8_1d: ... +def lanczos(M: _WindowLength, *, sym: op.CanBool = True) -> _Array_f8_1d: ... + +# Overloads where `return_ratios` is `False`. +@overload def dpss( - M: Untyped, - NW: Untyped, - Kmax: Untyped | None = None, - sym: bool = True, - norm: Untyped | None = None, - return_ratios: bool = False, -) -> Untyped: ... -def lanczos(M: Untyped, *, sym: bool = True) -> Untyped: ... -def get_window(window: Untyped, Nx: Untyped, fftbins: bool = True) -> Untyped: ... + M: _WindowLength, + NW: AnyReal, + Kmax: op.CanIndex, + sym: op.CanBool = True, + norm: _Norm | None = None, + return_ratios: Literal[False] = False, +) -> _Array_f8_2d: ... +@overload +def dpss( + M: _WindowLength, + NW: AnyReal, + Kmax: None = None, + sym: op.CanBool = True, + norm: _Norm | None = None, + return_ratios: Literal[False] = False, +) -> _Array_f8_1d: ... + +# Overloads where `return_ratios` is `True`. +# `return_ratios` as a positional argument +@overload +def dpss( + M: _WindowLength, + NW: AnyReal, + Kmax: op.CanIndex, + sym: op.CanBool, + norm: _Norm | None, + return_ratios: Literal[True], +) -> tuple[_Array_f8_2d, _Array_f8_1d]: ... + +# `return_ratios` as a keyword argument +@overload +def dpss( + M: _WindowLength, + NW: AnyReal, + Kmax: op.CanIndex, + sym: op.CanBool = True, + norm: _Norm | None = None, + *, + return_ratios: Literal[True], +) -> tuple[_Array_f8_2d, _Array_f8_1d]: ... + +# `return_ratios` as a positional argument +@overload +def dpss( + M: _WindowLength, + NW: AnyReal, + Kmax: None, + sym: op.CanBool, + norm: _Norm | None, + return_ratios: Literal[True], +) -> tuple[_Array_f8_1d, np.float64]: ... + +# `return_ratios` as a keyword argument +@overload +def dpss( + M: _WindowLength, + NW: AnyReal, + Kmax: None = None, + sym: op.CanBool = True, + norm: _Norm | None = None, + *, + return_ratios: Literal[True], +) -> tuple[_Array_f8_1d, np.float64]: ... + +# +def get_window( + window: _Window | AnyReal | tuple[_Window | _WindowNeedsParams, Unpack[tuple[object, ...]]], + Nx: _WindowLength, + fftbins: op.CanBool = True, +) -> _Array_f8_1d: ... diff --git a/tests/signal/windows/test_windows.pyi b/tests/signal/windows/test_windows.pyi new file mode 100644 index 00000000..7dbc5934 --- /dev/null +++ b/tests/signal/windows/test_windows.pyi @@ -0,0 +1,14 @@ +from typing import TypeAlias +from typing_extensions import assert_type + +import numpy as np +from scipy.signal.windows import dpss + +_Array_f8_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.float64]] +_Array_f8_2d: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.float64]] + +# test dpss function overloads +assert_type(dpss(64, 3), _Array_f8_1d) +assert_type(dpss(64, 3, 2), _Array_f8_2d) +assert_type(dpss(64, 3, return_ratios=True), tuple[_Array_f8_1d, np.float64]) +assert_type(dpss(64, 3, 2, return_ratios=True), tuple[_Array_f8_2d, _Array_f8_1d])