_helper.pyi 1.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  1. from typing import Any, Final, TypeVar, overload
  2. from typing import Literal as L
  3. from numpy import complexfloating, floating, generic, integer
  4. from numpy._typing import (
  5. ArrayLike,
  6. NDArray,
  7. _ArrayLike,
  8. _ArrayLikeComplex_co,
  9. _ArrayLikeFloat_co,
  10. _ShapeLike,
  11. )
  12. __all__ = ["fftfreq", "fftshift", "ifftshift", "rfftfreq"]
  13. _ScalarT = TypeVar("_ScalarT", bound=generic)
  14. ###
  15. integer_types: Final[tuple[type[int], type[integer]]] = ...
  16. ###
  17. @overload
  18. def fftshift(x: _ArrayLike[_ScalarT], axes: _ShapeLike | None = None) -> NDArray[_ScalarT]: ...
  19. @overload
  20. def fftshift(x: ArrayLike, axes: _ShapeLike | None = None) -> NDArray[Any]: ...
  21. #
  22. @overload
  23. def ifftshift(x: _ArrayLike[_ScalarT], axes: _ShapeLike | None = None) -> NDArray[_ScalarT]: ...
  24. @overload
  25. def ifftshift(x: ArrayLike, axes: _ShapeLike | None = None) -> NDArray[Any]: ...
  26. #
  27. @overload
  28. def fftfreq(n: int | integer, d: _ArrayLikeFloat_co = 1.0, device: L["cpu"] | None = None) -> NDArray[floating]: ...
  29. @overload
  30. def fftfreq(n: int | integer, d: _ArrayLikeComplex_co = 1.0, device: L["cpu"] | None = None) -> NDArray[complexfloating]: ...
  31. #
  32. @overload
  33. def rfftfreq(n: int | integer, d: _ArrayLikeFloat_co = 1.0, device: L["cpu"] | None = None) -> NDArray[floating]: ...
  34. @overload
  35. def rfftfreq(n: int | integer, d: _ArrayLikeComplex_co = 1.0, device: L["cpu"] | None = None) -> NDArray[complexfloating]: ...