Skip to content

Commit 5559ee0

Browse files
jorenhamjorenham
authored
Merge pull request #195 from pavyamsiri/improve/signal/waveforms
`signal`: Add type stubs to `_waveforms.pyi`.
2 parents d450251 + 76785f5 commit 5559ee0

File tree

2 files changed

+433
-22
lines changed

2 files changed

+433
-22
lines changed

‎scipy-stubs/signal/_waveforms.pyi‎

Lines changed: 225 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,229 @@
1-
from scipy._typing import Untyped
1+
from collections.abc import Iterable
2+
from typing import Any, Literal, TypeAlias, TypeVar, overload
3+
4+
import numpy as np
5+
import numpy.typing as npt
6+
import optype as op
7+
import optype.numpy as onp
8+
from numpy._typing import _DTypeLike
9+
from scipy._typing import AnyShape
210

311
__all__ = ["chirp", "gausspulse", "sawtooth", "square", "sweep_poly", "unit_impulse"]
412

5-
def sawtooth(t: Untyped, width: int = 1) -> Untyped: ...
6-
def square(t: Untyped, duty: float = 0.5) -> Untyped: ...
7-
def gausspulse(
8-
t: Untyped,
9-
fc: int = 1000,
10-
bw: float = 0.5,
11-
bwr: int = -6,
12-
tpr: int = -60,
13-
retquad: bool = False,
14-
retenv: bool = False,
15-
) -> Untyped: ...
13+
_SCT = TypeVar("_SCT", bound=np.generic)
14+
15+
_Truthy: TypeAlias = Literal[1, True]
16+
_Falsy: TypeAlias = Literal[0, False]
17+
_ArrayLikeFloat: TypeAlias = onp.ToFloat | onp.ToFloatND
18+
_Array_f: TypeAlias = onp.ArrayND[np.floating[Any]]
19+
_Array_f8: TypeAlias = onp.ArrayND[np.float64]
20+
21+
_ChirpMethod: TypeAlias = Literal["linear", "quadratic", "logarithmic", "hyperbolic"]
22+
23+
def sawtooth(t: _ArrayLikeFloat, width: _ArrayLikeFloat = 1) -> _Array_f8: ...
24+
def square(t: _ArrayLikeFloat, duty: _ArrayLikeFloat = 0.5) -> _Array_f8: ...
25+
26+
#
27+
@overload # Arrays
28+
def chirp(
29+
t: onp.ToFloatND,
30+
f0: onp.ToFloat,
31+
t1: onp.ToFloat,
32+
f1: onp.ToFloat,
33+
method: _ChirpMethod = "linear",
34+
phi: onp.ToFloat = 0,
35+
vertex_zero: op.CanBool = True,
36+
) -> _Array_f: ...
37+
@overload # Scalars
1638
def chirp(
17-
t: Untyped,
18-
f0: Untyped,
19-
t1: Untyped,
20-
f1: Untyped,
21-
method: str = "linear",
22-
phi: int = 0,
23-
vertex_zero: bool = True,
24-
) -> Untyped: ...
25-
def sweep_poly(t: Untyped, poly: Untyped, phi: int = 0) -> Untyped: ...
26-
def unit_impulse(shape: Untyped, idx: Untyped | None = None, dtype: Untyped = ...) -> Untyped: ...
39+
t: onp.ToFloat,
40+
f0: onp.ToFloat,
41+
t1: onp.ToFloat,
42+
f1: onp.ToFloat,
43+
method: _ChirpMethod = "linear",
44+
phi: onp.ToFloat = 0,
45+
vertex_zero: op.CanBool = True,
46+
) -> np.floating[Any]: ...
47+
48+
#
49+
def sweep_poly(
50+
t: _ArrayLikeFloat,
51+
poly: onp.ToFloatND | np.poly1d,
52+
phi: onp.ToFloat = 0,
53+
) -> _Array_f8: ...
54+
55+
#
56+
@overload # dtype is not given
57+
def unit_impulse(
58+
shape: AnyShape,
59+
idx: op.CanIndex | Iterable[op.CanIndex] | Literal["mid"] | None = None,
60+
dtype: type[float] = ...,
61+
) -> _Array_f8: ...
62+
@overload # dtype is given
63+
def unit_impulse(
64+
shape: AnyShape,
65+
idx: op.CanIndex | Iterable[op.CanIndex] | Literal["mid"] | None,
66+
dtype: _DTypeLike[_SCT],
67+
) -> npt.NDArray[_SCT]: ...
68+
69+
# Overloads for gausspulse when `t` is `"cutoff"`
70+
@overload # retquad: False = ..., retenv: False = ...
71+
def gausspulse(
72+
t: Literal["cutoff"],
73+
fc: onp.ToFloat = 1000,
74+
bw: onp.ToFloat = 0.5,
75+
bwr: onp.ToFloat = -6,
76+
tpr: onp.ToFloat = -60,
77+
retquad: op.CanBool = False,
78+
retenv: op.CanBool = False,
79+
) -> np.float64: ...
80+
81+
# Overloads for gausspulse when `t` is scalar
82+
@overload # retquad: False = ..., retenv: False = ...
83+
def gausspulse(
84+
t: onp.ToFloat,
85+
fc: onp.ToFloat = 1000,
86+
bw: onp.ToFloat = 0.5,
87+
bwr: onp.ToFloat = -6,
88+
tpr: onp.ToFloat = -60,
89+
retquad: _Falsy = False,
90+
retenv: _Falsy = False,
91+
) -> np.float64: ...
92+
@overload # retquad: False = ..., retenv: True (keyword)
93+
def gausspulse(
94+
t: onp.ToFloat,
95+
fc: onp.ToFloat = 1000,
96+
bw: onp.ToFloat = 0.5,
97+
bwr: onp.ToFloat = -6,
98+
tpr: onp.ToFloat = -60,
99+
retquad: _Falsy = False,
100+
*,
101+
retenv: _Truthy,
102+
) -> tuple[np.float64, np.float64]: ...
103+
@overload # retquad: False (positional), retenv: False (positional)
104+
def gausspulse(
105+
t: onp.ToFloat,
106+
fc: onp.ToFloat,
107+
bw: onp.ToFloat,
108+
bwr: onp.ToFloat,
109+
tpr: onp.ToFloat,
110+
retquad: _Falsy,
111+
retenv: _Truthy,
112+
) -> tuple[np.float64, np.float64]: ...
113+
@overload # retquad: True (positional), retenv: False = ...
114+
def gausspulse(
115+
t: onp.ToFloat,
116+
fc: onp.ToFloat,
117+
bw: onp.ToFloat,
118+
bwr: onp.ToFloat,
119+
tpr: onp.ToFloat,
120+
retquad: _Truthy,
121+
retenv: _Falsy = False,
122+
) -> tuple[np.float64, np.float64]: ...
123+
@overload # retquad: True (keyword), retenv: False = ...
124+
def gausspulse(
125+
t: onp.ToFloat,
126+
fc: onp.ToFloat = 1000,
127+
bw: onp.ToFloat = 0.5,
128+
bwr: onp.ToFloat = -6,
129+
tpr: onp.ToFloat = -60,
130+
*,
131+
retquad: _Truthy,
132+
retenv: _Falsy = False,
133+
) -> tuple[np.float64, np.float64]: ...
134+
@overload # retquad: True (positional), retenv: True (positional/keyword)
135+
def gausspulse(
136+
t: onp.ToFloat,
137+
fc: onp.ToFloat,
138+
bw: onp.ToFloat,
139+
bwr: onp.ToFloat,
140+
tpr: onp.ToFloat,
141+
retquad: _Truthy,
142+
retenv: _Truthy,
143+
) -> tuple[np.float64, np.float64, np.float64]: ...
144+
@overload # retquad: True (keyword), retenv: True
145+
def gausspulse(
146+
t: onp.ToFloat,
147+
fc: onp.ToFloat = 1000,
148+
bw: onp.ToFloat = 0.5,
149+
bwr: onp.ToFloat = -6,
150+
tpr: onp.ToFloat = -60,
151+
*,
152+
retquad: _Truthy,
153+
retenv: _Truthy,
154+
) -> tuple[np.float64, np.float64, np.float64]: ...
155+
156+
# Overloads for `gausspulse` when `t` is a non-scalar array like
157+
@overload # retquad: False = ..., retenv: False = ...
158+
def gausspulse(
159+
t: onp.ToFloatND,
160+
fc: onp.ToFloat = 1000,
161+
bw: onp.ToFloat = 0.5,
162+
bwr: onp.ToFloat = -6,
163+
tpr: onp.ToFloat = -60,
164+
retquad: _Falsy = False,
165+
retenv: _Falsy = False,
166+
) -> _Array_f8: ...
167+
@overload # retquad: False = ..., retenv: True (keyword)
168+
def gausspulse(
169+
t: onp.ToFloatND,
170+
fc: onp.ToFloat = 1000,
171+
bw: onp.ToFloat = 0.5,
172+
bwr: onp.ToFloat = -6,
173+
tpr: onp.ToFloat = -60,
174+
retquad: _Falsy = False,
175+
*,
176+
retenv: _Truthy,
177+
) -> tuple[_Array_f8, _Array_f8]: ...
178+
@overload # retquad: False (positional), retenv: False (positional)
179+
def gausspulse(
180+
t: onp.ToFloatND,
181+
fc: onp.ToFloat,
182+
bw: onp.ToFloat,
183+
bwr: onp.ToFloat,
184+
tpr: onp.ToFloat,
185+
retquad: _Falsy,
186+
retenv: _Truthy,
187+
) -> tuple[_Array_f8, _Array_f8]: ...
188+
@overload # retquad: True (positional), retenv: False = ...
189+
def gausspulse(
190+
t: onp.ToFloatND,
191+
fc: onp.ToFloat,
192+
bw: onp.ToFloat,
193+
bwr: onp.ToFloat,
194+
tpr: onp.ToFloat,
195+
retquad: _Truthy,
196+
retenv: _Falsy = False,
197+
) -> tuple[_Array_f8, _Array_f8]: ...
198+
@overload # retquad: True (keyword), retenv: False = ...
199+
def gausspulse(
200+
t: onp.ToFloatND,
201+
fc: onp.ToFloat = 1000,
202+
bw: onp.ToFloat = 0.5,
203+
bwr: onp.ToFloat = -6,
204+
tpr: onp.ToFloat = -60,
205+
*,
206+
retquad: _Truthy,
207+
retenv: _Falsy = False,
208+
) -> tuple[_Array_f8, _Array_f8]: ...
209+
@overload # retquad: True (positional), retenv: True (positional/keyword)
210+
def gausspulse(
211+
t: onp.ToFloatND,
212+
fc: onp.ToFloat,
213+
bw: onp.ToFloat,
214+
bwr: onp.ToFloat,
215+
tpr: onp.ToFloat,
216+
retquad: _Truthy,
217+
retenv: _Truthy,
218+
) -> tuple[_Array_f8, _Array_f8, _Array_f8]: ...
219+
@overload # retquad: True (keyword), retenv: True
220+
def gausspulse(
221+
t: onp.ToFloatND,
222+
fc: onp.ToFloat = 1000,
223+
bw: onp.ToFloat = 0.5,
224+
bwr: onp.ToFloat = -6,
225+
tpr: onp.ToFloat = -60,
226+
*,
227+
retquad: _Truthy,
228+
retenv: _Truthy,
229+
) -> tuple[_Array_f8, _Array_f8, _Array_f8]: ...

0 commit comments

Comments
 (0)