✨ slightly improved numpy.fft.[i]fftshift signatures

Author: jorenham

src/numpy-stubs/fft/_helper.pyi

@@ -1,33 +1,47 @@
-from typing import Any, Literal as L, overload
+from typing import Any, Final, Literal as L, TypeAlias, overload
 from typing_extensions import TypeVar
 
 import numpy as np
 from numpy._typing import ArrayLike, NDArray, _ArrayLike, _ArrayLikeComplex_co, _ArrayLikeFloat_co, _ShapeLike
 
 __all__ = ["fftfreq", "fftshift", "ifftshift", "rfftfreq"]
 
-_SCT = TypeVar("_SCT", bound=np.generic)
+###
+
+_ScalarT = TypeVar("_ScalarT", bound=np.generic)
+
+_Device: TypeAlias = L["cpu"] | None
+_IntLike: TypeAlias = int | np.integer
+
+###
+
+integer_types: Final[tuple[type[int], type[np.integer]]] = ...
+
+###
 
-#
 @overload
-def fftshift(x: _ArrayLike[_SCT], axes: _ShapeLike | None = ...) -> NDArray[_SCT]: ...
+def fftshift(x: _ArrayLike[_ScalarT], axes: _ShapeLike | None = None) -> NDArray[_ScalarT]: ...
 @overload
-def fftshift(x: ArrayLike, axes: _ShapeLike | None = ...) -> NDArray[Any]: ...
+def fftshift(x: _ArrayLikeFloat_co, axes: _ShapeLike | None = None) -> NDArray[np.floating]: ...
+@overload
+def fftshift(x: ArrayLike, axes: _ShapeLike | None = None) -> NDArray[Any]: ...
 
 #
 @overload
-def ifftshift(x: _ArrayLike[_SCT], axes: _ShapeLike | None = ...) -> NDArray[_SCT]: ...
+def ifftshift(x: _ArrayLike[_ScalarT], axes: _ShapeLike | None = None) -> NDArray[_ScalarT]: ...
+@overload
+def ifftshift(x: _ArrayLikeFloat_co, axes: _ShapeLike | None = None) -> NDArray[np.floating]: ...
 @overload
-def ifftshift(x: ArrayLike, axes: _ShapeLike | None = ...) -> NDArray[Any]: ...
+def ifftshift(x: ArrayLike, axes: _ShapeLike | None = None) -> NDArray[Any]: ...
 
 #
 @overload
-def fftfreq(n: int | np.integer, d: _ArrayLikeFloat_co = ..., device: L["cpu"] | None = ...) -> NDArray[np.floating]: ...
+def fftfreq(n: _IntLike, d: _ArrayLikeFloat_co = 1.0, device: _Device = None) -> NDArray[np.floating]: ...
 @overload
-def fftfreq(n: int | np.integer, d: _ArrayLikeComplex_co = ..., device: L["cpu"] | None = ...) -> NDArray[np.inexact]: ...
+def fftfreq(n: _IntLike, d: _ArrayLikeComplex_co = 1.0, device: _Device = None) -> NDArray[np.inexact]: ...
 
 #
 @overload
-def rfftfreq(n: int | np.integer, d: _ArrayLikeFloat_co = ..., device: L["cpu"] | None = ...) -> NDArray[np.floating]: ...
+def rfftfreq(n: _IntLike, d: _ArrayLikeFloat_co = 1.0, device: _Device = None) -> NDArray[np.floating]: ...
 @overload
-def rfftfreq(n: int | np.integer, d: _ArrayLikeComplex_co = ..., device: L["cpu"] | None = ...) -> NDArray[np.inexact]: ...
+def rfftfreq(n: _IntLike, d: _ArrayLikeComplex_co = 1.0, device: _Device = None) -> NDArray[np.inexact]: ...

test/static/accept/fft.pyi

@@ -1,4 +1,3 @@
-from typing import Any
 from typing_extensions import assert_type
 
 import numpy as np
@@ -9,10 +8,10 @@ AR_c16: npt.NDArray[np.complex128]
 AR_LIKE_f8: list[float]
 
 assert_type(np.fft.fftshift(AR_f8), npt.NDArray[np.float64])
-assert_type(np.fft.fftshift(AR_LIKE_f8, axes=0), npt.NDArray[Any])
+assert_type(np.fft.fftshift(AR_LIKE_f8, axes=0), npt.NDArray[np.floating])
 
 assert_type(np.fft.ifftshift(AR_f8), npt.NDArray[np.float64])
-assert_type(np.fft.ifftshift(AR_LIKE_f8, axes=0), npt.NDArray[Any])
+assert_type(np.fft.ifftshift(AR_LIKE_f8, axes=0), npt.NDArray[np.floating])
 
 assert_type(np.fft.fftfreq(5, AR_f8), npt.NDArray[np.floating])
 assert_type(np.fft.fftfreq(np.int64(), AR_c16), npt.NDArray[np.inexact])