Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

signal: Add type stubs to _waveforms.pyi. #195

Merged
merged 19 commits into from
Nov 25, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
36e770f
`signal`: Add type stubs for `sawtooth` and `square` in `_waveforms.p…
pavyamsiri Nov 23, 2024
157b713
`signal`: Add type stubs for `sweep_poly` in `_waveforms.pyi`
pavyamsiri Nov 23, 2024
810224f
`signal`: Add type stubs to `unit_impulse` in `_waveforms.pyi`
pavyamsiri Nov 23, 2024
dbe336e
`signal`+tests: Add type stubs for `gausspulse` in `_waveforms.pyi`.
pavyamsiri Nov 23, 2024
7a63f34
`signal`: Add type stubs to `chirp` in `_waveforms.pyi`.
pavyamsiri Nov 23, 2024
7c8aa77
`signal`: Don't allow onp.ToComplexND because time should be real-valued
pavyamsiri Nov 23, 2024
d416367
`signal`: Use a type alias for `npt.NDArray[np.float64]`
pavyamsiri Nov 23, 2024
cfb7713
`signal`: Sort imports in `_waveforms.pyi`
pavyamsiri Nov 23, 2024
2d60ea0
`signal`: Omit default `float` value in `_waveforms`.
pavyamsiri Nov 23, 2024
f2a870b
`signal`: Fix up some mistakes in `_waveforms.pyi`.
pavyamsiri Nov 24, 2024
bb7ec09
`signal`: Allow scalars as `t` in `gausspulse`
pavyamsiri Nov 24, 2024
d656dd6
tests: Update `test_waveforms`
pavyamsiri Nov 24, 2024
7524442
`signal`: `t` in `gausspulse` accepts the string "cutoff" not the `tpr`
pavyamsiri Nov 24, 2024
7b975ec
tests: Update `test_waveforms`
pavyamsiri Nov 24, 2024
647b1c3
`signal`: Type annotate `chirp` so that the output dtype can be inferred
pavyamsiri Nov 24, 2024
3278292
`signal`+tests: Add all possible overloads for `gausspulse`.
pavyamsiri Nov 24, 2024
4d74dc8
`signal`: Move `gausspulse` overloads around to satisfy mypy
pavyamsiri Nov 24, 2024
60570a0
`signal`: Fix `chirp` overload overlap in `_waveforms.pyi`
pavyamsiri Nov 25, 2024
76785f5
`signal`: Simplify overloads for `chirp` in `_waveforms.pyi`
pavyamsiri Nov 25, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
247 changes: 225 additions & 22 deletions scipy-stubs/signal/_waveforms.pyi
Original file line number Diff line number Diff line change
@@ -1,26 +1,229 @@
from scipy._typing import Untyped
from collections.abc import Iterable
from typing import Any, Literal, TypeAlias, TypeVar, overload

import numpy as np
import numpy.typing as npt
import optype as op
import optype.numpy as onp
from numpy._typing import _DTypeLike
from scipy._typing import AnyShape

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

def sawtooth(t: Untyped, width: int = 1) -> Untyped: ...
def square(t: Untyped, duty: float = 0.5) -> Untyped: ...
def gausspulse(
t: Untyped,
fc: int = 1000,
bw: float = 0.5,
bwr: int = -6,
tpr: int = -60,
retquad: bool = False,
retenv: bool = False,
) -> Untyped: ...
_SCT = TypeVar("_SCT", bound=np.generic)

_Truthy: TypeAlias = Literal[1, True]
_Falsy: TypeAlias = Literal[0, False]
_ArrayLikeFloat: TypeAlias = onp.ToFloat | onp.ToFloatND
_Array_f: TypeAlias = onp.ArrayND[np.floating[Any]]
_Array_f8: TypeAlias = onp.ArrayND[np.float64]

_ChirpMethod: TypeAlias = Literal["linear", "quadratic", "logarithmic", "hyperbolic"]

def sawtooth(t: _ArrayLikeFloat, width: _ArrayLikeFloat = 1) -> _Array_f8: ...
def square(t: _ArrayLikeFloat, duty: _ArrayLikeFloat = 0.5) -> _Array_f8: ...

#
@overload # Arrays
def chirp(
t: onp.ToFloatND,
f0: onp.ToFloat,
t1: onp.ToFloat,
f1: onp.ToFloat,
method: _ChirpMethod = "linear",
phi: onp.ToFloat = 0,
vertex_zero: op.CanBool = True,
) -> _Array_f: ...
@overload # Scalars
def chirp(
t: Untyped,
f0: Untyped,
t1: Untyped,
f1: Untyped,
method: str = "linear",
phi: int = 0,
vertex_zero: bool = True,
) -> Untyped: ...
def sweep_poly(t: Untyped, poly: Untyped, phi: int = 0) -> Untyped: ...
def unit_impulse(shape: Untyped, idx: Untyped | None = None, dtype: Untyped = ...) -> Untyped: ...
t: onp.ToFloat,
f0: onp.ToFloat,
t1: onp.ToFloat,
f1: onp.ToFloat,
method: _ChirpMethod = "linear",
phi: onp.ToFloat = 0,
vertex_zero: op.CanBool = True,
) -> np.floating[Any]: ...

#
def sweep_poly(
t: _ArrayLikeFloat,
poly: onp.ToFloatND | np.poly1d,
phi: onp.ToFloat = 0,
) -> _Array_f8: ...

#
@overload # dtype is not given
def unit_impulse(
shape: AnyShape,
idx: op.CanIndex | Iterable[op.CanIndex] | Literal["mid"] | None = None,
dtype: type[float] = ...,
) -> _Array_f8: ...
@overload # dtype is given
def unit_impulse(
shape: AnyShape,
idx: op.CanIndex | Iterable[op.CanIndex] | Literal["mid"] | None,
dtype: _DTypeLike[_SCT],
) -> npt.NDArray[_SCT]: ...

# Overloads for gausspulse when `t` is `"cutoff"`
@overload # retquad: False = ..., retenv: False = ...
def gausspulse(
t: Literal["cutoff"],
fc: onp.ToFloat = 1000,
bw: onp.ToFloat = 0.5,
bwr: onp.ToFloat = -6,
tpr: onp.ToFloat = -60,
retquad: op.CanBool = False,
retenv: op.CanBool = False,
) -> np.float64: ...

# Overloads for gausspulse when `t` is scalar
@overload # retquad: False = ..., retenv: False = ...
def gausspulse(
t: onp.ToFloat,
fc: onp.ToFloat = 1000,
bw: onp.ToFloat = 0.5,
bwr: onp.ToFloat = -6,
tpr: onp.ToFloat = -60,
retquad: _Falsy = False,
retenv: _Falsy = False,
) -> np.float64: ...
@overload # retquad: False = ..., retenv: True (keyword)
def gausspulse(
t: onp.ToFloat,
fc: onp.ToFloat = 1000,
bw: onp.ToFloat = 0.5,
bwr: onp.ToFloat = -6,
tpr: onp.ToFloat = -60,
retquad: _Falsy = False,
*,
retenv: _Truthy,
) -> tuple[np.float64, np.float64]: ...
@overload # retquad: False (positional), retenv: False (positional)
def gausspulse(
t: onp.ToFloat,
fc: onp.ToFloat,
bw: onp.ToFloat,
bwr: onp.ToFloat,
tpr: onp.ToFloat,
retquad: _Falsy,
retenv: _Truthy,
) -> tuple[np.float64, np.float64]: ...
@overload # retquad: True (positional), retenv: False = ...
def gausspulse(
t: onp.ToFloat,
fc: onp.ToFloat,
bw: onp.ToFloat,
bwr: onp.ToFloat,
tpr: onp.ToFloat,
retquad: _Truthy,
retenv: _Falsy = False,
) -> tuple[np.float64, np.float64]: ...
@overload # retquad: True (keyword), retenv: False = ...
def gausspulse(
t: onp.ToFloat,
fc: onp.ToFloat = 1000,
bw: onp.ToFloat = 0.5,
bwr: onp.ToFloat = -6,
tpr: onp.ToFloat = -60,
*,
retquad: _Truthy,
retenv: _Falsy = False,
) -> tuple[np.float64, np.float64]: ...
@overload # retquad: True (positional), retenv: True (positional/keyword)
def gausspulse(
t: onp.ToFloat,
fc: onp.ToFloat,
bw: onp.ToFloat,
bwr: onp.ToFloat,
tpr: onp.ToFloat,
retquad: _Truthy,
retenv: _Truthy,
) -> tuple[np.float64, np.float64, np.float64]: ...
@overload # retquad: True (keyword), retenv: True
def gausspulse(
t: onp.ToFloat,
fc: onp.ToFloat = 1000,
bw: onp.ToFloat = 0.5,
bwr: onp.ToFloat = -6,
tpr: onp.ToFloat = -60,
*,
retquad: _Truthy,
retenv: _Truthy,
) -> tuple[np.float64, np.float64, np.float64]: ...

# Overloads for `gausspulse` when `t` is a non-scalar array like
@overload # retquad: False = ..., retenv: False = ...
def gausspulse(
t: onp.ToFloatND,
fc: onp.ToFloat = 1000,
bw: onp.ToFloat = 0.5,
bwr: onp.ToFloat = -6,
tpr: onp.ToFloat = -60,
retquad: _Falsy = False,
retenv: _Falsy = False,
) -> _Array_f8: ...
@overload # retquad: False = ..., retenv: True (keyword)
def gausspulse(
t: onp.ToFloatND,
fc: onp.ToFloat = 1000,
bw: onp.ToFloat = 0.5,
bwr: onp.ToFloat = -6,
tpr: onp.ToFloat = -60,
retquad: _Falsy = False,
*,
retenv: _Truthy,
) -> tuple[_Array_f8, _Array_f8]: ...
@overload # retquad: False (positional), retenv: False (positional)
def gausspulse(
t: onp.ToFloatND,
fc: onp.ToFloat,
bw: onp.ToFloat,
bwr: onp.ToFloat,
tpr: onp.ToFloat,
retquad: _Falsy,
retenv: _Truthy,
) -> tuple[_Array_f8, _Array_f8]: ...
@overload # retquad: True (positional), retenv: False = ...
def gausspulse(
t: onp.ToFloatND,
fc: onp.ToFloat,
bw: onp.ToFloat,
bwr: onp.ToFloat,
tpr: onp.ToFloat,
retquad: _Truthy,
retenv: _Falsy = False,
) -> tuple[_Array_f8, _Array_f8]: ...
@overload # retquad: True (keyword), retenv: False = ...
def gausspulse(
t: onp.ToFloatND,
fc: onp.ToFloat = 1000,
bw: onp.ToFloat = 0.5,
bwr: onp.ToFloat = -6,
tpr: onp.ToFloat = -60,
*,
retquad: _Truthy,
retenv: _Falsy = False,
) -> tuple[_Array_f8, _Array_f8]: ...
@overload # retquad: True (positional), retenv: True (positional/keyword)
def gausspulse(
t: onp.ToFloatND,
fc: onp.ToFloat,
bw: onp.ToFloat,
bwr: onp.ToFloat,
tpr: onp.ToFloat,
retquad: _Truthy,
retenv: _Truthy,
) -> tuple[_Array_f8, _Array_f8, _Array_f8]: ...
@overload # retquad: True (keyword), retenv: True
def gausspulse(
t: onp.ToFloatND,
fc: onp.ToFloat = 1000,
bw: onp.ToFloat = 0.5,
bwr: onp.ToFloat = -6,
tpr: onp.ToFloat = -60,
*,
retquad: _Truthy,
retenv: _Truthy,
) -> tuple[_Array_f8, _Array_f8, _Array_f8]: ...
Loading