Skip to content

Commit 818150c

Browse files
jorenhamjorenham
authored
Merge pull request #153 from pavyamsiri/improve/windows
`signal.windows`: Add type stubs for `_windows.pyi`.
2 parents 58a090d + 8cc7c34 commit 818150c

File tree

2 files changed

+198
-33
lines changed

2 files changed

+198
-33
lines changed

‎scipy-stubs/signal/windows/_windows.pyi

+184-33
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,11 @@
1-
from scipy._typing import Untyped
1+
from collections.abc import Sequence
2+
from typing import Any, Literal, TypeAlias, overload
3+
from typing_extensions import Unpack
4+
5+
import numpy as np
6+
import numpy.typing as npt
7+
import optype as op
8+
from scipy._typing import AnyInt, AnyReal
29

310
__all__ = [
411
"barthann",
@@ -29,36 +36,180 @@ __all__ = [
2936
"tukey",
3037
]
3138

32-
def general_cosine(M: Untyped, a: Untyped, sym: bool = True) -> Untyped: ...
33-
def boxcar(M: Untyped, sym: bool = True) -> Untyped: ...
34-
def triang(M: Untyped, sym: bool = True) -> Untyped: ...
35-
def parzen(M: Untyped, sym: bool = True) -> Untyped: ...
36-
def bohman(M: Untyped, sym: bool = True) -> Untyped: ...
37-
def blackman(M: Untyped, sym: bool = True) -> Untyped: ...
38-
def nuttall(M: Untyped, sym: bool = True) -> Untyped: ...
39-
def blackmanharris(M: Untyped, sym: bool = True) -> Untyped: ...
40-
def flattop(M: Untyped, sym: bool = True) -> Untyped: ...
41-
def bartlett(M: Untyped, sym: bool = True) -> Untyped: ...
42-
def hann(M: Untyped, sym: bool = True) -> Untyped: ...
43-
def tukey(M: Untyped, alpha: float = 0.5, sym: bool = True) -> Untyped: ...
44-
def barthann(M: Untyped, sym: bool = True) -> Untyped: ...
45-
def general_hamming(M: Untyped, alpha: Untyped, sym: bool = True) -> Untyped: ...
46-
def hamming(M: Untyped, sym: bool = True) -> Untyped: ...
47-
def kaiser(M: Untyped, beta: Untyped, sym: bool = True) -> Untyped: ...
48-
def kaiser_bessel_derived(M: Untyped, beta: Untyped, *, sym: bool = True) -> Untyped: ...
49-
def gaussian(M: Untyped, std: Untyped, sym: bool = True) -> Untyped: ...
50-
def general_gaussian(M: Untyped, p: Untyped, sig: Untyped, sym: bool = True) -> Untyped: ...
51-
def chebwin(M: Untyped, at: Untyped, sym: bool = True) -> Untyped: ...
52-
def cosine(M: Untyped, sym: bool = True) -> Untyped: ...
53-
def exponential(M: Untyped, center: Untyped | None = None, tau: float = 1.0, sym: bool = True) -> Untyped: ...
54-
def taylor(M: Untyped, nbar: int = 4, sll: int = 30, norm: bool = True, sym: bool = True) -> Untyped: ...
39+
_Array_f8_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.float64]]
40+
_Array_f8_2d: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.float64]]
41+
_Weights: TypeAlias = Sequence[AnyReal] | npt.NDArray[np.floating[Any]] | npt.NDArray[np.integer[Any]] | npt.NDArray[np.bool_]
42+
43+
_Norm: TypeAlias = Literal[2, "approximate", "subsample"]
44+
_WindowLength: TypeAlias = int | np.int16 | np.int32 | np.int64
45+
_Window: TypeAlias = Literal[
46+
"barthann",
47+
"brthan",
48+
"bth",
49+
"bartlett",
50+
"bart",
51+
"brt",
52+
"blackman",
53+
"black",
54+
"blk",
55+
"blackmanharris",
56+
"blackharr",
57+
"bkh",
58+
"bohman",
59+
"bman",
60+
"bmn",
61+
"boxcar",
62+
"box",
63+
"ones",
64+
"rect",
65+
"rectangular",
66+
"cosine",
67+
"halfcosine",
68+
"exponential",
69+
"poisson",
70+
"flattop",
71+
"flat",
72+
"flt",
73+
"hamming",
74+
"hamm",
75+
"ham",
76+
"hann",
77+
"han",
78+
"lanczos",
79+
"sinc",
80+
"nuttall",
81+
"nutl",
82+
"nut",
83+
"parzen",
84+
"parz",
85+
"par",
86+
"taylor",
87+
"taylorwin",
88+
"triangle",
89+
"triang",
90+
"tri",
91+
"tukey",
92+
"tuk",
93+
]
94+
_WindowNeedsParams: TypeAlias = Literal[
95+
"chebwin",
96+
"cheb",
97+
"dpss",
98+
"gaussian",
99+
"gauss",
100+
"gss",
101+
"general cosine",
102+
"general_cosine",
103+
"general gaussian",
104+
"general_gaussian",
105+
"general gauss",
106+
"general_gauss",
107+
"ggs",
108+
"general hamming",
109+
"general_hamming",
110+
"kaiser",
111+
"ksr",
112+
"kaiser bessel derived",
113+
"kbd",
114+
]
115+
116+
def general_cosine(M: _WindowLength, a: _Weights, sym: op.CanBool = True) -> _Array_f8_1d: ...
117+
def boxcar(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ...
118+
def triang(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ...
119+
def parzen(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ...
120+
def bohman(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ...
121+
def blackman(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ...
122+
def nuttall(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ...
123+
def blackmanharris(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ...
124+
def flattop(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ...
125+
def bartlett(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ...
126+
def hann(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ...
127+
def tukey(M: _WindowLength, alpha: AnyReal = 0.5, sym: op.CanBool = True) -> _Array_f8_1d: ...
128+
def barthann(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ...
129+
def general_hamming(M: _WindowLength, alpha: AnyReal, sym: op.CanBool = True) -> _Array_f8_1d: ...
130+
def hamming(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ...
131+
def kaiser(M: _WindowLength, beta: AnyReal, sym: op.CanBool = True) -> _Array_f8_1d: ...
132+
def kaiser_bessel_derived(M: _WindowLength, beta: AnyReal, *, sym: op.CanBool = True) -> _Array_f8_1d: ...
133+
def gaussian(M: _WindowLength, std: AnyReal, sym: op.CanBool = True) -> _Array_f8_1d: ...
134+
def general_gaussian(M: _WindowLength, p: AnyReal, sig: AnyReal, sym: op.CanBool = True) -> _Array_f8_1d: ...
135+
def chebwin(M: _WindowLength, at: AnyReal, sym: op.CanBool = True) -> _Array_f8_1d: ...
136+
def cosine(M: _WindowLength, sym: op.CanBool = True) -> _Array_f8_1d: ...
137+
def exponential(M: _WindowLength, center: AnyReal | None = None, tau: AnyReal = 1.0, sym: op.CanBool = True) -> _Array_f8_1d: ...
138+
def taylor(
139+
M: _WindowLength, nbar: AnyInt = 4, sll: AnyInt = 30, norm: op.CanBool = True, sym: op.CanBool = True
140+
) -> _Array_f8_1d: ...
141+
def lanczos(M: _WindowLength, *, sym: op.CanBool = True) -> _Array_f8_1d: ...
142+
143+
# Overloads where `return_ratios` is `False`.
144+
@overload
55145
def dpss(
56-
M: Untyped,
57-
NW: Untyped,
58-
Kmax: Untyped | None = None,
59-
sym: bool = True,
60-
norm: Untyped | None = None,
61-
return_ratios: bool = False,
62-
) -> Untyped: ...
63-
def lanczos(M: Untyped, *, sym: bool = True) -> Untyped: ...
64-
def get_window(window: Untyped, Nx: Untyped, fftbins: bool = True) -> Untyped: ...
146+
M: _WindowLength,
147+
NW: AnyReal,
148+
Kmax: op.CanIndex,
149+
sym: op.CanBool = True,
150+
norm: _Norm | None = None,
151+
return_ratios: Literal[False] = False,
152+
) -> _Array_f8_2d: ...
153+
@overload
154+
def dpss(
155+
M: _WindowLength,
156+
NW: AnyReal,
157+
Kmax: None = None,
158+
sym: op.CanBool = True,
159+
norm: _Norm | None = None,
160+
return_ratios: Literal[False] = False,
161+
) -> _Array_f8_1d: ...
162+
163+
# Overloads where `return_ratios` is `True`.
164+
# `return_ratios` as a positional argument
165+
@overload
166+
def dpss(
167+
M: _WindowLength,
168+
NW: AnyReal,
169+
Kmax: op.CanIndex,
170+
sym: op.CanBool,
171+
norm: _Norm | None,
172+
return_ratios: Literal[True],
173+
) -> tuple[_Array_f8_2d, _Array_f8_1d]: ...
174+
175+
# `return_ratios` as a keyword argument
176+
@overload
177+
def dpss(
178+
M: _WindowLength,
179+
NW: AnyReal,
180+
Kmax: op.CanIndex,
181+
sym: op.CanBool = True,
182+
norm: _Norm | None = None,
183+
*,
184+
return_ratios: Literal[True],
185+
) -> tuple[_Array_f8_2d, _Array_f8_1d]: ...
186+
187+
# `return_ratios` as a positional argument
188+
@overload
189+
def dpss(
190+
M: _WindowLength,
191+
NW: AnyReal,
192+
Kmax: None,
193+
sym: op.CanBool,
194+
norm: _Norm | None,
195+
return_ratios: Literal[True],
196+
) -> tuple[_Array_f8_1d, np.float64]: ...
197+
198+
# `return_ratios` as a keyword argument
199+
@overload
200+
def dpss(
201+
M: _WindowLength,
202+
NW: AnyReal,
203+
Kmax: None = None,
204+
sym: op.CanBool = True,
205+
norm: _Norm | None = None,
206+
*,
207+
return_ratios: Literal[True],
208+
) -> tuple[_Array_f8_1d, np.float64]: ...
209+
210+
#
211+
def get_window(
212+
window: _Window | AnyReal | tuple[_Window | _WindowNeedsParams, Unpack[tuple[object, ...]]],
213+
Nx: _WindowLength,
214+
fftbins: op.CanBool = True,
215+
) -> _Array_f8_1d: ...

‎tests/signal/windows/test_windows.pyi

+14
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
from typing import TypeAlias
2+
from typing_extensions import assert_type
3+
4+
import numpy as np
5+
from scipy.signal.windows import dpss
6+
7+
_Array_f8_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.float64]]
8+
_Array_f8_2d: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.float64]]
9+
10+
# test dpss function overloads
11+
assert_type(dpss(64, 3), _Array_f8_1d)
12+
assert_type(dpss(64, 3, 2), _Array_f8_2d)
13+
assert_type(dpss(64, 3, return_ratios=True), tuple[_Array_f8_1d, np.float64])
14+
assert_type(dpss(64, 3, 2, return_ratios=True), tuple[_Array_f8_2d, _Array_f8_1d])

0 commit comments

Comments
 (0)