Skip to content

Commit 1d9b2ac

Browse files
committed
ndimage: type-tests for the fourier_* functions
1 parent 7541bdc commit 1d9b2ac

File tree

1 file changed

+93
-0
lines changed

1 file changed

+93
-0
lines changed

tests/ndimage/test__fourier.pyi

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
# type-tests for `ndimage/_fourier.pyi`
2+
3+
from typing import TypeAlias, assert_type
4+
5+
import numpy as np
6+
import numpy.typing as npt
7+
import optype.numpy as onp
8+
9+
from scipy.ndimage import fourier_gaussian, fourier_shift
10+
11+
i8_nd: npt.NDArray[np.int8]
12+
f16_nd: npt.NDArray[np.float16]
13+
f32_nd: npt.NDArray[np.float32]
14+
f64_nd: npt.NDArray[np.float64]
15+
c64_nd: npt.NDArray[np.complex64]
16+
c128_nd: npt.NDArray[np.complex128]
17+
18+
int_2d: list[list[int]]
19+
float_2d: list[list[float]]
20+
complex_2d: list[list[complex]]
21+
22+
_OutputArray: TypeAlias = onp.Array2D[np.complex64]
23+
_OutputArrayND: TypeAlias = onp.ArrayND[np.complex64]
24+
output_array: _OutputArray
25+
output_sctype: type[np.complex64]
26+
27+
###
28+
# `fourier_gaussian` (also `fourier_ellipsoid` and `fourier_uniform`)
29+
# NOTE: `fourier_uniform` and `fourier_ellipsoid` have the same signature, so no need to also test those.
30+
31+
assert_type(fourier_gaussian(i8_nd, 4), onp.ArrayND[np.float64])
32+
assert_type(fourier_gaussian(f16_nd, 4), onp.ArrayND[np.float32])
33+
assert_type(fourier_gaussian(f32_nd, 4), onp.ArrayND[np.float32])
34+
assert_type(fourier_gaussian(f64_nd, 4), onp.ArrayND[np.float64])
35+
assert_type(fourier_gaussian(c64_nd, 4), onp.ArrayND[np.complex64])
36+
assert_type(fourier_gaussian(c128_nd, 4), onp.ArrayND[np.complex128])
37+
assert_type(fourier_gaussian(int_2d, 4), onp.ArrayND[np.float64])
38+
assert_type(fourier_gaussian(float_2d, 4), onp.ArrayND[np.float64])
39+
assert_type(fourier_gaussian(complex_2d, 4), onp.ArrayND[np.complex128])
40+
41+
assert_type(fourier_gaussian(i8_nd, 4, output=output_array), _OutputArray)
42+
assert_type(fourier_gaussian(f16_nd, 4, output=output_array), _OutputArray)
43+
assert_type(fourier_gaussian(f32_nd, 4, output=output_array), _OutputArray)
44+
assert_type(fourier_gaussian(f64_nd, 4, output=output_array), _OutputArray)
45+
assert_type(fourier_gaussian(c64_nd, 4, output=output_array), _OutputArray)
46+
assert_type(fourier_gaussian(c128_nd, 4, output=output_array), _OutputArray)
47+
assert_type(fourier_gaussian(int_2d, 4, output=output_array), _OutputArray)
48+
assert_type(fourier_gaussian(float_2d, 4, output=output_array), _OutputArray)
49+
assert_type(fourier_gaussian(complex_2d, 4, output=output_array), _OutputArray)
50+
51+
assert_type(fourier_gaussian(i8_nd, 4, output=output_sctype), _OutputArrayND)
52+
assert_type(fourier_gaussian(f16_nd, 4, output=output_sctype), _OutputArrayND)
53+
assert_type(fourier_gaussian(f32_nd, 4, output=output_sctype), _OutputArrayND)
54+
assert_type(fourier_gaussian(f64_nd, 4, output=output_sctype), _OutputArrayND)
55+
assert_type(fourier_gaussian(c64_nd, 4, output=output_sctype), _OutputArrayND)
56+
assert_type(fourier_gaussian(c128_nd, 4, output=output_sctype), _OutputArrayND)
57+
assert_type(fourier_gaussian(int_2d, 4, output=output_sctype), _OutputArrayND)
58+
assert_type(fourier_gaussian(float_2d, 4, output=output_sctype), _OutputArrayND)
59+
assert_type(fourier_gaussian(complex_2d, 4, output=output_sctype), _OutputArrayND)
60+
61+
###
62+
# `fourier_shift`
63+
# NOTE: Unlike the other three functions, this always returns complex output.
64+
65+
assert_type(fourier_shift(i8_nd, 4), onp.ArrayND[np.complex128])
66+
assert_type(fourier_shift(f16_nd, 4), onp.ArrayND[np.complex128])
67+
assert_type(fourier_shift(f32_nd, 4), onp.ArrayND[np.complex128])
68+
assert_type(fourier_shift(f64_nd, 4), onp.ArrayND[np.complex128])
69+
assert_type(fourier_shift(c64_nd, 4), onp.ArrayND[np.complex64])
70+
assert_type(fourier_shift(c128_nd, 4), onp.ArrayND[np.complex128])
71+
assert_type(fourier_shift(int_2d, 4), onp.ArrayND[np.complex128])
72+
assert_type(fourier_shift(float_2d, 4), onp.ArrayND[np.complex128])
73+
assert_type(fourier_shift(complex_2d, 4), onp.ArrayND[np.complex128])
74+
75+
assert_type(fourier_shift(i8_nd, 4, output=output_array), _OutputArray)
76+
assert_type(fourier_shift(f16_nd, 4, output=output_array), _OutputArray)
77+
assert_type(fourier_shift(f32_nd, 4, output=output_array), _OutputArray)
78+
assert_type(fourier_shift(f64_nd, 4, output=output_array), _OutputArray)
79+
assert_type(fourier_shift(c64_nd, 4, output=output_array), _OutputArray)
80+
assert_type(fourier_shift(c128_nd, 4, output=output_array), _OutputArray)
81+
assert_type(fourier_shift(int_2d, 4, output=output_array), _OutputArray)
82+
assert_type(fourier_shift(float_2d, 4, output=output_array), _OutputArray)
83+
assert_type(fourier_shift(complex_2d, 4, output=output_array), _OutputArray)
84+
85+
assert_type(fourier_shift(i8_nd, 4, output=output_sctype), _OutputArrayND)
86+
assert_type(fourier_shift(f16_nd, 4, output=output_sctype), _OutputArrayND)
87+
assert_type(fourier_shift(f32_nd, 4, output=output_sctype), _OutputArrayND)
88+
assert_type(fourier_shift(f64_nd, 4, output=output_sctype), _OutputArrayND)
89+
assert_type(fourier_shift(c64_nd, 4, output=output_sctype), _OutputArrayND)
90+
assert_type(fourier_shift(c128_nd, 4, output=output_sctype), _OutputArrayND)
91+
assert_type(fourier_shift(int_2d, 4, output=output_sctype), _OutputArrayND)
92+
assert_type(fourier_shift(float_2d, 4, output=output_sctype), _OutputArrayND)
93+
assert_type(fourier_shift(complex_2d, 4, output=output_sctype), _OutputArrayND)

0 commit comments

Comments
 (0)