diff --git a/scipy-stubs/signal/_waveforms.pyi b/scipy-stubs/signal/_waveforms.pyi index 11389919..4886184b 100644 --- a/scipy-stubs/signal/_waveforms.pyi +++ b/scipy-stubs/signal/_waveforms.pyi @@ -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]: ... diff --git a/tests/signal/test_waveforms.pyi b/tests/signal/test_waveforms.pyi new file mode 100644 index 00000000..94bbf253 --- /dev/null +++ b/tests/signal/test_waveforms.pyi @@ -0,0 +1,208 @@ +from typing import Literal, TypeAlias +from typing_extensions import assert_type + +import numpy as np +import optype.numpy as onp +from scipy.signal import gausspulse + +_Array_f8: TypeAlias = onp.ArrayND[np.float64] +_Scalar_f8: TypeAlias = np.float64 +_Falsy: TypeAlias = Literal[0, False] +_Truthy: TypeAlias = Literal[1, True] + +_time_scalar: onp.ToFloat +_time_array: onp.ToFloatND +_time_cutoff: Literal["cutoff"] +_float: onp.ToFloat +_truthy: _Truthy +_falsy: _Falsy + +# test gausspulse function overloads +# `time` as an array +assert_type(gausspulse(_time_array), _Array_f8) +# Full positional arguments +assert_type(gausspulse(_time_array, _float, _float, _float, _float, _falsy, _falsy), _Array_f8) +assert_type( + gausspulse(_time_array, _float, _float, _float, _float, _truthy, _falsy), + tuple[_Array_f8, _Array_f8], +) +assert_type( + gausspulse(_time_array, _float, _float, _float, _float, _falsy, _truthy), + tuple[_Array_f8, _Array_f8], +) +assert_type( + gausspulse(_time_array, _float, _float, _float, _float, _truthy, _truthy), + tuple[_Array_f8, _Array_f8, _Array_f8], +) +# Full keyword arguments +assert_type(gausspulse(t=_time_array), _Array_f8) +assert_type( + gausspulse( + t=_time_array, + retquad=_falsy, + retenv=_falsy, + ), + _Array_f8, +) +assert_type( + gausspulse( + t=_time_array, + retquad=_truthy, + retenv=_falsy, + ), + tuple[_Array_f8, _Array_f8], +) +assert_type( + gausspulse( + t=_time_array, + retquad=_falsy, + retenv=_truthy, + ), + tuple[_Array_f8, _Array_f8], +) +assert_type( + gausspulse( + t=_time_array, + retquad=_truthy, + retenv=_truthy, + ), + tuple[_Array_f8, _Array_f8, _Array_f8], +) + +# Mixed positional and keyword arguments +assert_type( + gausspulse(_time_array, _float, _float, _float, _float, retquad=_falsy, retenv=_falsy), + _Array_f8, +) +assert_type( + gausspulse(_time_array, _float, _float, _float, _float, retquad=_truthy, retenv=_falsy), + tuple[_Array_f8, _Array_f8], +) +assert_type( + gausspulse(_time_array, _float, _float, _float, _float, retquad=_falsy, retenv=_truthy), + tuple[_Array_f8, _Array_f8], +) +assert_type( + gausspulse(_time_array, _float, _float, _float, _float, retquad=_truthy, retenv=_truthy), + tuple[_Array_f8, _Array_f8, _Array_f8], +) + +# `time` as a scalar +assert_type(gausspulse(_time_scalar), _Scalar_f8) +# Full positional arguments +assert_type(gausspulse(_time_scalar, _float, _float, _float, _float, _falsy, _falsy), _Scalar_f8) +assert_type( + gausspulse(_time_scalar, _float, _float, _float, _float, _truthy, _falsy), + tuple[_Scalar_f8, _Scalar_f8], +) +assert_type( + gausspulse(_time_scalar, _float, _float, _float, _float, _falsy, _truthy), + tuple[_Scalar_f8, _Scalar_f8], +) +assert_type( + gausspulse(_time_scalar, _float, _float, _float, _float, _truthy, _truthy), + tuple[_Scalar_f8, _Scalar_f8, _Scalar_f8], +) +# Full keyword arguments +assert_type(gausspulse(t=_time_scalar), _Scalar_f8) +assert_type( + gausspulse( + t=_time_scalar, + retquad=_falsy, + retenv=_falsy, + ), + _Scalar_f8, +) +assert_type( + gausspulse( + t=_time_scalar, + retquad=_truthy, + retenv=_falsy, + ), + tuple[_Scalar_f8, _Scalar_f8], +) +assert_type( + gausspulse( + t=_time_scalar, + retquad=_falsy, + retenv=_truthy, + ), + tuple[_Scalar_f8, _Scalar_f8], +) +assert_type( + gausspulse( + t=_time_scalar, + retquad=_truthy, + retenv=_truthy, + ), + tuple[_Scalar_f8, _Scalar_f8, _Scalar_f8], +) + +# Mixed positional and keyword arguments +assert_type( + gausspulse(_time_scalar, _float, _float, _float, _float, retquad=_falsy, retenv=_falsy), + _Scalar_f8, +) +assert_type( + gausspulse(_time_scalar, _float, _float, _float, _float, retquad=_truthy, retenv=_falsy), + tuple[_Scalar_f8, _Scalar_f8], +) +assert_type( + gausspulse(_time_scalar, _float, _float, _float, _float, retquad=_falsy, retenv=_truthy), + tuple[_Scalar_f8, _Scalar_f8], +) +assert_type( + gausspulse(_time_scalar, _float, _float, _float, _float, retquad=_truthy, retenv=_truthy), + tuple[_Scalar_f8, _Scalar_f8, _Scalar_f8], +) + +# `time` as the literal `"cutoff"` +assert_type(gausspulse(_time_cutoff), _Scalar_f8) +# Full positional arguments +assert_type(gausspulse(_time_cutoff, _float, _float, _float, _float, _falsy, _falsy), _Scalar_f8) +assert_type(gausspulse(_time_cutoff, _float, _float, _float, _float, _truthy, _falsy), _Scalar_f8) +assert_type(gausspulse(_time_cutoff, _float, _float, _float, _float, _falsy, _truthy), _Scalar_f8) +assert_type(gausspulse(_time_cutoff, _float, _float, _float, _float, _truthy, _truthy), _Scalar_f8) +# Full keyword arguments +assert_type(gausspulse(t=_time_cutoff), _Scalar_f8) +assert_type( + gausspulse( + t=_time_cutoff, + retquad=_falsy, + retenv=_falsy, + ), + _Scalar_f8, +) +assert_type( + gausspulse( + t=_time_cutoff, + retquad=_truthy, + retenv=_falsy, + ), + _Scalar_f8, +) +assert_type( + gausspulse( + t=_time_cutoff, + retquad=_falsy, + retenv=_truthy, + ), + _Scalar_f8, +) +assert_type( + gausspulse( + t=_time_cutoff, + retquad=_truthy, + retenv=_truthy, + ), + _Scalar_f8, +) + +# Mixed positional and keyword arguments +assert_type( + gausspulse(_time_cutoff, _float, _float, _float, _float, retquad=_falsy, retenv=_falsy), + _Scalar_f8, +) +assert_type(gausspulse(_time_cutoff, _float, _float, _float, _float, retquad=_truthy, retenv=_falsy), _Scalar_f8) +assert_type(gausspulse(_time_cutoff, _float, _float, _float, _float, retquad=_falsy, retenv=_truthy), _Scalar_f8) +assert_type(gausspulse(_time_cutoff, _float, _float, _float, _float, retquad=_truthy, retenv=_truthy), _Scalar_f8)