signal.windows: Add type stubs for _windows.pyi.

scipy-stubs/signal/windows/_windows.pyi

@@ -1,4 +1,10 @@
-from scipy._typing import Untyped
+from collections.abc import Sequence
+from typing import Literal, TypeAlias, overload
+from typing_extensions import Unpack
+
+import numpy as np
+import optype as op
+from scipy._typing import AnyInt, AnyReal
 
 __all__ = [
     "barthann",
@@ -29,36 +35,179 @@ __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]]
+
+_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: Sequence[AnyReal], 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: ...

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])