|
| 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