Skip to content

Commit 55aba94

Browse files
jorenhamjorenham
committed
🎨 refactor from _numtype import in fft 
1 parent d3bd20f commit 55aba94

1 file changed

Lines changed: 73 additions & 88 deletions

File tree

‎src/numpy-stubs/fft/_helper.pyi‎

Lines changed: 73 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -1,29 +1,8 @@
11
from typing import Any, Final, Literal as L, TypeAlias, overload
22
from typing_extensions import TypeVar
33

4+
import _numtype as _nt
45
import numpy as np
5-
from _numtype import (
6-
Array,
7-
Array1D,
8-
CanLenArray,
9-
JustComplex,
10-
ToBool_1d,
11-
ToBool_1nd,
12-
ToBytes_1d,
13-
ToBytes_1nd,
14-
ToComplex128_1d,
15-
ToComplex128_1nd,
16-
ToFloat64_1d,
17-
ToFloat64_1nd,
18-
ToGeneric_1nd,
19-
ToInt_1d,
20-
ToInt_1nd,
21-
ToStr_1d,
22-
ToStr_1nd,
23-
co_complex,
24-
co_float,
25-
co_integer,
26-
)
276
from numpy._typing import _ShapeLike
287

298
__all__ = ["fftfreq", "fftshift", "ifftshift", "rfftfreq"]
@@ -44,140 +23,146 @@ integer_types: Final[tuple[type[int], type[np.integer]]] = ...
4423

4524
# keep in sync with `ifftshift`
4625
@overload
47-
def fftshift(x: ToBool_1d, axes: _ShapeLike | None = None) -> Array1D[np.bool]: ...
26+
def fftshift(x: _nt.ToBool_1d, axes: _ShapeLike | None = None) -> _nt.Array1D[np.bool]: ...
4827
@overload
49-
def fftshift(x: ToBool_1nd, axes: _ShapeLike | None = None) -> Array[np.bool]: ...
28+
def fftshift(x: _nt.ToBool_1nd, axes: _ShapeLike | None = None) -> _nt.Array[np.bool]: ...
5029
@overload
51-
def fftshift(x: CanLenArray[_ScalarT, _ShapeT], axes: _ShapeLike | None = None) -> Array[_ScalarT, _ShapeT]: ...
30+
def fftshift(x: _nt.CanLenArray[_ScalarT, _ShapeT], axes: _ShapeLike | None = None) -> _nt.Array[_ScalarT, _ShapeT]: ...
5231
@overload
53-
def fftshift(x: ToInt_1d, axes: _ShapeLike | None = None) -> Array1D[np.intp]: ...
32+
def fftshift(x: _nt.ToInt_1d, axes: _ShapeLike | None = None) -> _nt.Array1D[np.intp]: ...
5433
@overload
55-
def fftshift(x: ToInt_1nd, axes: _ShapeLike | None = None) -> Array[np.intp]: ...
34+
def fftshift(x: _nt.ToInt_1nd, axes: _ShapeLike | None = None) -> _nt.Array[np.intp]: ...
5635
@overload
57-
def fftshift(x: ToFloat64_1d, axes: _ShapeLike | None = None) -> Array1D[np.float64]: ...
36+
def fftshift(x: _nt.ToFloat64_1d, axes: _ShapeLike | None = None) -> _nt.Array1D[np.float64]: ...
5837
@overload
59-
def fftshift(x: ToFloat64_1nd, axes: _ShapeLike | None = None) -> Array[np.float64]: ...
38+
def fftshift(x: _nt.ToFloat64_1nd, axes: _ShapeLike | None = None) -> _nt.Array[np.float64]: ...
6039
@overload
61-
def fftshift(x: ToComplex128_1d, axes: _ShapeLike | None = None) -> Array1D[np.complex128]: ...
40+
def fftshift(x: _nt.ToComplex128_1d, axes: _ShapeLike | None = None) -> _nt.Array1D[np.complex128]: ...
6241
@overload
63-
def fftshift(x: ToComplex128_1nd, axes: _ShapeLike | None = None) -> Array[np.complex128]: ...
42+
def fftshift(x: _nt.ToComplex128_1nd, axes: _ShapeLike | None = None) -> _nt.Array[np.complex128]: ...
6443
@overload
65-
def fftshift(x: ToBytes_1d, axes: _ShapeLike | None = None) -> Array1D[np.bytes_]: ...
44+
def fftshift(x: _nt.ToBytes_1d, axes: _ShapeLike | None = None) -> _nt.Array1D[np.bytes_]: ...
6645
@overload
67-
def fftshift(x: ToBytes_1nd, axes: _ShapeLike | None = None) -> Array[np.bytes_]: ...
46+
def fftshift(x: _nt.ToBytes_1nd, axes: _ShapeLike | None = None) -> _nt.Array[np.bytes_]: ...
6847
@overload
69-
def fftshift(x: ToStr_1d, axes: _ShapeLike | None = None) -> Array1D[np.str_]: ...
48+
def fftshift(x: _nt.ToStr_1d, axes: _ShapeLike | None = None) -> _nt.Array1D[np.str_]: ...
7049
@overload
71-
def fftshift(x: ToStr_1nd, axes: _ShapeLike | None = None) -> Array[np.str_]: ...
50+
def fftshift(x: _nt.ToStr_1nd, axes: _ShapeLike | None = None) -> _nt.Array[np.str_]: ...
7251
@overload
73-
def fftshift(x: ToGeneric_1nd, axes: _ShapeLike | None = None) -> Array[Any]: ...
52+
def fftshift(x: _nt.ToGeneric_1nd, axes: _ShapeLike | None = None) -> _nt.Array[Any]: ...
7453

7554
# keep in sync with `fftshift`
7655
@overload
77-
def ifftshift(x: ToBool_1d, axes: _ShapeLike | None = None) -> Array1D[np.bool]: ...
56+
def ifftshift(x: _nt.ToBool_1d, axes: _ShapeLike | None = None) -> _nt.Array1D[np.bool]: ...
7857
@overload
79-
def ifftshift(x: ToBool_1nd, axes: _ShapeLike | None = None) -> Array[np.bool]: ...
58+
def ifftshift(x: _nt.ToBool_1nd, axes: _ShapeLike | None = None) -> _nt.Array[np.bool]: ...
8059
@overload
81-
def ifftshift(x: CanLenArray[_ScalarT, _ShapeT], axes: _ShapeLike | None = None) -> Array[_ScalarT, _ShapeT]: ...
60+
def ifftshift(
61+
x: _nt.CanLenArray[_ScalarT, _ShapeT], axes: _ShapeLike | None = None
62+
) -> _nt.Array[_ScalarT, _ShapeT]: ...
8263
@overload
83-
def ifftshift(x: ToInt_1d, axes: _ShapeLike | None = None) -> Array1D[np.intp]: ...
64+
def ifftshift(x: _nt.ToInt_1d, axes: _ShapeLike | None = None) -> _nt.Array1D[np.intp]: ...
8465
@overload
85-
def ifftshift(x: ToInt_1nd, axes: _ShapeLike | None = None) -> Array[np.intp]: ...
66+
def ifftshift(x: _nt.ToInt_1nd, axes: _ShapeLike | None = None) -> _nt.Array[np.intp]: ...
8667
@overload
87-
def ifftshift(x: ToFloat64_1d, axes: _ShapeLike | None = None) -> Array1D[np.float64]: ...
68+
def ifftshift(x: _nt.ToFloat64_1d, axes: _ShapeLike | None = None) -> _nt.Array1D[np.float64]: ...
8869
@overload
89-
def ifftshift(x: ToFloat64_1nd, axes: _ShapeLike | None = None) -> Array[np.float64]: ...
70+
def ifftshift(x: _nt.ToFloat64_1nd, axes: _ShapeLike | None = None) -> _nt.Array[np.float64]: ...
9071
@overload
91-
def ifftshift(x: ToComplex128_1d, axes: _ShapeLike | None = None) -> Array1D[np.complex128]: ...
72+
def ifftshift(x: _nt.ToComplex128_1d, axes: _ShapeLike | None = None) -> _nt.Array1D[np.complex128]: ...
9273
@overload
93-
def ifftshift(x: ToComplex128_1nd, axes: _ShapeLike | None = None) -> Array[np.complex128]: ...
74+
def ifftshift(x: _nt.ToComplex128_1nd, axes: _ShapeLike | None = None) -> _nt.Array[np.complex128]: ...
9475
@overload
95-
def ifftshift(x: ToBytes_1d, axes: _ShapeLike | None = None) -> Array1D[np.bytes_]: ...
76+
def ifftshift(x: _nt.ToBytes_1d, axes: _ShapeLike | None = None) -> _nt.Array1D[np.bytes_]: ...
9677
@overload
97-
def ifftshift(x: ToBytes_1nd, axes: _ShapeLike | None = None) -> Array[np.bytes_]: ...
78+
def ifftshift(x: _nt.ToBytes_1nd, axes: _ShapeLike | None = None) -> _nt.Array[np.bytes_]: ...
9879
@overload
99-
def ifftshift(x: ToStr_1d, axes: _ShapeLike | None = None) -> Array1D[np.str_]: ...
80+
def ifftshift(x: _nt.ToStr_1d, axes: _ShapeLike | None = None) -> _nt.Array1D[np.str_]: ...
10081
@overload
101-
def ifftshift(x: ToStr_1nd, axes: _ShapeLike | None = None) -> Array[np.str_]: ...
82+
def ifftshift(x: _nt.ToStr_1nd, axes: _ShapeLike | None = None) -> _nt.Array[np.str_]: ...
10283
@overload
103-
def ifftshift(x: ToGeneric_1nd, axes: _ShapeLike | None = None) -> Array[Any]: ...
84+
def ifftshift(x: _nt.ToGeneric_1nd, axes: _ShapeLike | None = None) -> _nt.Array[Any]: ...
10485

10586
# keep in sync with `rfftfreq`
10687
@overload # 0d float | integer -> 1d float64
107-
def fftfreq(n: _Int, d: float | co_integer = 1.0, device: _Device | None = None) -> Array1D[np.float64]: ...
88+
def fftfreq(n: _Int, d: float | _nt.co_integer = 1.0, device: _Device | None = None) -> _nt.Array1D[np.float64]: ...
10889
@overload # 0d longdouble -> 1d longdouble
109-
def fftfreq(n: _Int, d: np.longdouble, device: _Device | None = None) -> Array1D[np.longdouble]: ...
90+
def fftfreq(n: _Int, d: np.longdouble, device: _Device | None = None) -> _nt.Array1D[np.longdouble]: ...
11091
@overload # 0d float{16,32,64} -> 1d float64
111-
def fftfreq(n: _Int, d: _Floating64Max, device: _Device | None = None) -> Array1D[np.float64]: ...
92+
def fftfreq(n: _Int, d: _Floating64Max, device: _Device | None = None) -> _nt.Array1D[np.float64]: ...
11293
@overload # 0d complex -> 1d complex128
113-
def fftfreq(n: _Int, d: JustComplex, device: _Device | None = None) -> Array1D[np.complex128]: ...
94+
def fftfreq(n: _Int, d: _nt.JustComplex, device: _Device | None = None) -> _nt.Array1D[np.complex128]: ...
11495
@overload # 0d clongdouble -> 1d clongdouble
115-
def fftfreq(n: _Int, d: np.clongdouble, device: _Device | None = None) -> Array1D[np.clongdouble]: ...
96+
def fftfreq(n: _Int, d: np.clongdouble, device: _Device | None = None) -> _nt.Array1D[np.clongdouble]: ...
11697
@overload # 0d complex{64,128} -> 1d complex128
117-
def fftfreq(n: _Int, d: _CFloating64Max, device: _Device | None = None) -> Array1D[np.complex128]: ...
98+
def fftfreq(n: _Int, d: _CFloating64Max, device: _Device | None = None) -> _nt.Array1D[np.complex128]: ...
11899
@overload # Nd +integer -> Nd float64
119-
def fftfreq(n: _Int, d: Array[co_integer, _ShapeT], device: _Device | None = None) -> Array[np.float64, _ShapeT]: ...
100+
def fftfreq(
101+
n: _Int, d: _nt.Array[_nt.co_integer, _ShapeT], device: _Device | None = None
102+
) -> _nt.Array[np.float64, _ShapeT]: ...
120103
@overload # Nd longdouble -> Nd longdouble
121104
def fftfreq(
122-
n: _Int, d: Array[np.longdouble, _ShapeT], device: _Device | None = None
123-
) -> Array[np.longdouble, _ShapeT]: ...
105+
n: _Int, d: _nt.Array[np.longdouble, _ShapeT], device: _Device | None = None
106+
) -> _nt.Array[np.longdouble, _ShapeT]: ...
124107
@overload # Nd float{16,32,64} -> Nd float64
125108
def fftfreq(
126-
n: _Int, d: Array[_Floating64Max, _ShapeT], device: _Device | None = None
127-
) -> Array[np.float64, _ShapeT]: ...
109+
n: _Int, d: _nt.Array[_Floating64Max, _ShapeT], device: _Device | None = None
110+
) -> _nt.Array[np.float64, _ShapeT]: ...
128111
@overload # Nd complex longdouble -> Nd complex longdouble
129112
def fftfreq(
130-
n: _Int, d: Array[np.clongdouble, _ShapeT], device: _Device | None = None
131-
) -> Array[np.clongdouble, _ShapeT]: ...
113+
n: _Int, d: _nt.Array[np.clongdouble, _ShapeT], device: _Device | None = None
114+
) -> _nt.Array[np.clongdouble, _ShapeT]: ...
132115
@overload # Nd complex{64,128} -> Nd complex128
133116
def fftfreq(
134-
n: _Int, d: Array[_CFloating64Max, _ShapeT], device: _Device | None = None
135-
) -> Array[np.complex128, _ShapeT]: ...
117+
n: _Int, d: _nt.Array[_CFloating64Max, _ShapeT], device: _Device | None = None
118+
) -> _nt.Array[np.complex128, _ShapeT]: ...
136119
@overload # ?d +floating -> Nd floating
137120
def fftfreq(
138-
n: _Int, d: float | co_float | Array[co_float] = 1.0, device: _Device | None = None
139-
) -> Array[np.floating]: ...
121+
n: _Int, d: float | _nt.co_float | _nt.Array[_nt.co_float] = 1.0, device: _Device | None = None
122+
) -> _nt.Array[np.floating]: ...
140123
@overload # ?d +complex -> Nd inexact
141124
def fftfreq(
142-
n: _Int, d: complex | co_complex | Array[co_complex] = 1.0, device: _Device | None = None
143-
) -> Array[np.inexact]: ...
125+
n: _Int, d: complex | _nt.co_complex | _nt.Array[_nt.co_complex] = 1.0, device: _Device | None = None
126+
) -> _nt.Array[np.inexact]: ...
144127

145128
# keep in sync with `fftfreq`
146129
@overload # 0d float | integer -> 1d float64
147-
def rfftfreq(n: _Int, d: float | co_integer = 1.0, device: _Device | None = None) -> Array1D[np.float64]: ...
130+
def rfftfreq(n: _Int, d: float | _nt.co_integer = 1.0, device: _Device | None = None) -> _nt.Array1D[np.float64]: ...
148131
@overload # 0d longdouble -> 1d longdouble
149-
def rfftfreq(n: _Int, d: np.longdouble, device: _Device | None = None) -> Array1D[np.longdouble]: ...
132+
def rfftfreq(n: _Int, d: np.longdouble, device: _Device | None = None) -> _nt.Array1D[np.longdouble]: ...
150133
@overload # 0d float{16,32,64} -> 1d float64
151-
def rfftfreq(n: _Int, d: _Floating64Max, device: _Device | None = None) -> Array1D[np.float64]: ...
134+
def rfftfreq(n: _Int, d: _Floating64Max, device: _Device | None = None) -> _nt.Array1D[np.float64]: ...
152135
@overload # 0d complex -> 1d complex128
153-
def rfftfreq(n: _Int, d: JustComplex, device: _Device | None = None) -> Array1D[np.complex128]: ...
136+
def rfftfreq(n: _Int, d: _nt.JustComplex, device: _Device | None = None) -> _nt.Array1D[np.complex128]: ...
154137
@overload # 0d clongdouble -> 1d clongdouble
155-
def rfftfreq(n: _Int, d: np.clongdouble, device: _Device | None = None) -> Array1D[np.clongdouble]: ...
138+
def rfftfreq(n: _Int, d: np.clongdouble, device: _Device | None = None) -> _nt.Array1D[np.clongdouble]: ...
156139
@overload # 0d complex{64,128} -> 1d complex128
157-
def rfftfreq(n: _Int, d: _CFloating64Max, device: _Device | None = None) -> Array1D[np.complex128]: ...
140+
def rfftfreq(n: _Int, d: _CFloating64Max, device: _Device | None = None) -> _nt.Array1D[np.complex128]: ...
158141
@overload # Nd +integer -> Nd float64
159-
def rfftfreq(n: _Int, d: Array[co_integer, _ShapeT], device: _Device | None = None) -> Array[np.float64, _ShapeT]: ...
142+
def rfftfreq(
143+
n: _Int, d: _nt.Array[_nt.co_integer, _ShapeT], device: _Device | None = None
144+
) -> _nt.Array[np.float64, _ShapeT]: ...
160145
@overload # Nd longdouble -> Nd longdouble
161146
def rfftfreq(
162-
n: _Int, d: Array[np.longdouble, _ShapeT], device: _Device | None = None
163-
) -> Array[np.longdouble, _ShapeT]: ...
147+
n: _Int, d: _nt.Array[np.longdouble, _ShapeT], device: _Device | None = None
148+
) -> _nt.Array[np.longdouble, _ShapeT]: ...
164149
@overload # Nd float{16,32,64} -> Nd float64
165150
def rfftfreq(
166-
n: _Int, d: Array[_Floating64Max, _ShapeT], device: _Device | None = None
167-
) -> Array[np.float64, _ShapeT]: ...
151+
n: _Int, d: _nt.Array[_Floating64Max, _ShapeT], device: _Device | None = None
152+
) -> _nt.Array[np.float64, _ShapeT]: ...
168153
@overload # Nd complex longdouble -> Nd complex longdouble
169154
def rfftfreq(
170-
n: _Int, d: Array[np.clongdouble, _ShapeT], device: _Device | None = None
171-
) -> Array[np.clongdouble, _ShapeT]: ...
155+
n: _Int, d: _nt.Array[np.clongdouble, _ShapeT], device: _Device | None = None
156+
) -> _nt.Array[np.clongdouble, _ShapeT]: ...
172157
@overload # Nd complex{64,128} -> Nd complex128
173158
def rfftfreq(
174-
n: _Int, d: Array[_CFloating64Max, _ShapeT], device: _Device | None = None
175-
) -> Array[np.complex128, _ShapeT]: ...
159+
n: _Int, d: _nt.Array[_CFloating64Max, _ShapeT], device: _Device | None = None
160+
) -> _nt.Array[np.complex128, _ShapeT]: ...
176161
@overload # ?d +floating -> Nd floating
177162
def rfftfreq(
178-
n: _Int, d: float | co_float | Array[co_float] = 1.0, device: _Device | None = None
179-
) -> Array[np.floating]: ...
163+
n: _Int, d: float | _nt.co_float | _nt.Array[_nt.co_float] = 1.0, device: _Device | None = None
164+
) -> _nt.Array[np.floating]: ...
180165
@overload # ?d +complex -> Nd inexact
181166
def rfftfreq(
182-
n: _Int, d: complex | co_complex | Array[co_complex] = 1.0, device: _Device | None = None
183-
) -> Array[np.inexact]: ...
167+
n: _Int, d: complex | _nt.co_complex | _nt.Array[_nt.co_complex] = 1.0, device: _Device | None = None
168+
) -> _nt.Array[np.inexact]: ...

0 commit comments

Comments
 (0)