|
1 | 1 | from typing import Any |
2 | 2 | from typing_extensions import assert_type |
3 | 3 |
|
| 4 | +import _numtype as _nt |
4 | 5 | import numpy as np |
5 | | -import numpy.typing as npt |
6 | | -from _numtype import Array1D |
7 | 6 | from numpy.lib._arraysetops_impl import UniqueAllResult, UniqueCountsResult, UniqueInverseResult |
8 | 7 |
|
9 | | -AR_b: npt.NDArray[np.bool] |
10 | | -AR_i_: npt.NDArray[np.intp] |
11 | | -AR_f8: npt.NDArray[np.float64] |
12 | | -AR_M: npt.NDArray[np.datetime64] |
13 | | -AR_O: npt.NDArray[np.object_] |
| 8 | +AR_b: _nt.Array[np.bool] |
| 9 | +AR_i_: _nt.Array[np.intp] |
| 10 | +AR_f8: _nt.Array[np.float64] |
| 11 | +AR_M: _nt.Array[np.datetime64] |
| 12 | +AR_O: _nt.Array[np.object_] |
14 | 13 | AR_LIKE_i_: list[int] |
15 | 14 |
|
16 | 15 | ### |
17 | 16 |
|
18 | | -assert_type(np.intersect1d(AR_i_, AR_i_), Array1D[np.intp]) |
19 | | -assert_type(np.intersect1d(AR_f8, AR_i_), Array1D[np.float64]) |
20 | | -assert_type(np.intersect1d(AR_M, AR_M, assume_unique=True), Array1D[np.datetime64]) |
| 17 | +assert_type(np.intersect1d(AR_i_, AR_i_), _nt.Array1D[np.intp]) |
| 18 | +assert_type(np.intersect1d(AR_f8, AR_i_), _nt.Array1D[np.float64]) |
| 19 | +assert_type(np.intersect1d(AR_M, AR_M, assume_unique=True), _nt.Array1D[np.datetime64]) |
21 | 20 | assert_type( |
22 | 21 | np.intersect1d(AR_f8, AR_f8, return_indices=True), |
23 | | - tuple[Array1D[np.float64], Array1D[np.intp], Array1D[np.intp]], |
| 22 | + tuple[_nt.Array1D[np.float64], _nt.Array1D[np.intp], _nt.Array1D[np.intp]], |
24 | 23 | ) |
25 | 24 |
|
26 | | -assert_type(np.union1d(AR_i_, AR_i_), Array1D[np.intp]) |
27 | | -assert_type(np.union1d(AR_f8, AR_i_), Array1D[np.float64]) |
28 | | -assert_type(np.union1d(AR_M, AR_M), Array1D[np.datetime64]) |
| 25 | +assert_type(np.union1d(AR_i_, AR_i_), _nt.Array1D[np.intp]) |
| 26 | +assert_type(np.union1d(AR_f8, AR_i_), _nt.Array1D[np.float64]) |
| 27 | +assert_type(np.union1d(AR_M, AR_M), _nt.Array1D[np.datetime64]) |
29 | 28 |
|
30 | | -assert_type(np.ediff1d(AR_b), Array1D[np.int8]) |
31 | | -assert_type(np.ediff1d(AR_M), Array1D[np.timedelta64]) |
32 | | -assert_type(np.ediff1d(AR_O), Array1D[np.object_]) |
33 | | -assert_type(np.ediff1d(AR_i_, to_end=[1, 2, 3]), Array1D[np.intp]) |
34 | | -assert_type(np.ediff1d(AR_LIKE_i_, to_begin=[0, 1]), Array1D[np.intp]) |
| 29 | +assert_type(np.ediff1d(AR_b), _nt.Array1D[np.int8]) |
| 30 | +assert_type(np.ediff1d(AR_M), _nt.Array1D[np.timedelta64]) |
| 31 | +assert_type(np.ediff1d(AR_O), _nt.Array1D[np.object_]) |
| 32 | +assert_type(np.ediff1d(AR_i_, to_end=[1, 2, 3]), _nt.Array1D[np.intp]) |
| 33 | +assert_type(np.ediff1d(AR_LIKE_i_, to_begin=[0, 1]), _nt.Array1D[np.intp]) |
35 | 34 |
|
36 | | -assert_type(np.setxor1d(AR_i_, AR_i_), Array1D[np.intp]) |
37 | | -assert_type(np.setxor1d(AR_f8, AR_i_), Array1D[np.float64]) |
38 | | -assert_type(np.setxor1d(AR_M, AR_M, assume_unique=True), Array1D[np.datetime64]) |
| 35 | +assert_type(np.setxor1d(AR_i_, AR_i_), _nt.Array1D[np.intp]) |
| 36 | +assert_type(np.setxor1d(AR_f8, AR_i_), _nt.Array1D[np.float64]) |
| 37 | +assert_type(np.setxor1d(AR_M, AR_M, assume_unique=True), _nt.Array1D[np.datetime64]) |
39 | 38 |
|
40 | | -assert_type(np.setdiff1d(AR_i_, AR_i_), Array1D[np.intp]) |
41 | | -assert_type(np.setdiff1d(AR_f8, AR_i_), Array1D[np.float64]) |
42 | | -assert_type(np.setdiff1d(AR_M, AR_M, assume_unique=True), Array1D[np.datetime64]) |
| 39 | +assert_type(np.setdiff1d(AR_i_, AR_i_), _nt.Array1D[np.intp]) |
| 40 | +assert_type(np.setdiff1d(AR_f8, AR_i_), _nt.Array1D[np.float64]) |
| 41 | +assert_type(np.setdiff1d(AR_M, AR_M, assume_unique=True), _nt.Array1D[np.datetime64]) |
43 | 42 |
|
44 | | -assert_type(np.isin(AR_i_, AR_i_), npt.NDArray[np.bool]) |
45 | | -assert_type(np.isin(AR_f8, AR_i_), npt.NDArray[np.bool]) |
46 | | -assert_type(np.isin(AR_M, AR_M, assume_unique=True), npt.NDArray[np.bool]) |
47 | | -assert_type(np.isin(AR_f8, AR_LIKE_i_, invert=True), npt.NDArray[np.bool]) |
| 43 | +assert_type(np.isin(AR_i_, AR_i_), _nt.Array[np.bool]) |
| 44 | +assert_type(np.isin(AR_f8, AR_i_), _nt.Array[np.bool]) |
| 45 | +assert_type(np.isin(AR_M, AR_M, assume_unique=True), _nt.Array[np.bool]) |
| 46 | +assert_type(np.isin(AR_f8, AR_LIKE_i_, invert=True), _nt.Array[np.bool]) |
48 | 47 |
|
49 | | -assert_type(np.unique(AR_f8), npt.NDArray[np.float64]) |
50 | | -assert_type(np.unique(AR_f8, return_index=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp]]) |
51 | | -assert_type(np.unique(AR_f8, return_inverse=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp]]) |
52 | | -assert_type(np.unique(AR_f8, return_counts=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp]]) |
53 | | -assert_type(np.unique(AR_LIKE_i_, axis=0), npt.NDArray[Any]) |
54 | | -assert_type(np.unique(AR_LIKE_i_, return_index=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp]]) |
55 | | -assert_type(np.unique(AR_LIKE_i_, return_inverse=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp]]) |
56 | | -assert_type(np.unique(AR_LIKE_i_, return_counts=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp]]) |
| 48 | +assert_type(np.unique(AR_f8), _nt.Array[np.float64]) |
| 49 | +assert_type(np.unique(AR_f8, return_index=True), tuple[_nt.Array[np.float64], _nt.Array[np.intp]]) |
| 50 | +assert_type(np.unique(AR_f8, return_inverse=True), tuple[_nt.Array[np.float64], _nt.Array[np.intp]]) |
| 51 | +assert_type(np.unique(AR_f8, return_counts=True), tuple[_nt.Array[np.float64], _nt.Array[np.intp]]) |
| 52 | +assert_type(np.unique(AR_LIKE_i_, axis=0), _nt.Array[Any]) |
| 53 | +assert_type(np.unique(AR_LIKE_i_, return_index=True), tuple[_nt.Array[Any], _nt.Array[np.intp]]) |
| 54 | +assert_type(np.unique(AR_LIKE_i_, return_inverse=True), tuple[_nt.Array[Any], _nt.Array[np.intp]]) |
| 55 | +assert_type(np.unique(AR_LIKE_i_, return_counts=True), tuple[_nt.Array[Any], _nt.Array[np.intp]]) |
57 | 56 | assert_type( |
58 | 57 | np.unique(AR_f8, return_index=True, return_inverse=True), |
59 | | - tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp]], |
| 58 | + tuple[_nt.Array[np.float64], _nt.Array[np.intp], _nt.Array[np.intp]], |
60 | 59 | ) |
61 | 60 | assert_type( |
62 | 61 | np.unique(AR_LIKE_i_, return_index=True, return_inverse=True), |
63 | | - tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp]], |
| 62 | + tuple[_nt.Array[Any], _nt.Array[np.intp], _nt.Array[np.intp]], |
64 | 63 | ) |
65 | 64 | assert_type( |
66 | 65 | np.unique(AR_f8, return_index=True, return_counts=True), |
67 | | - tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp]], |
| 66 | + tuple[_nt.Array[np.float64], _nt.Array[np.intp], _nt.Array[np.intp]], |
68 | 67 | ) |
69 | 68 | assert_type( |
70 | 69 | np.unique(AR_LIKE_i_, return_index=True, return_counts=True), |
71 | | - tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp]], |
| 70 | + tuple[_nt.Array[Any], _nt.Array[np.intp], _nt.Array[np.intp]], |
72 | 71 | ) |
73 | 72 | assert_type( |
74 | 73 | np.unique(AR_f8, return_inverse=True, return_counts=True), |
75 | | - tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp]], |
| 74 | + tuple[_nt.Array[np.float64], _nt.Array[np.intp], _nt.Array[np.intp]], |
76 | 75 | ) |
77 | 76 | assert_type( |
78 | 77 | np.unique(AR_LIKE_i_, return_inverse=True, return_counts=True), |
79 | | - tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp]], |
| 78 | + tuple[_nt.Array[Any], _nt.Array[np.intp], _nt.Array[np.intp]], |
80 | 79 | ) |
81 | 80 | assert_type( |
82 | 81 | np.unique(AR_f8, return_index=True, return_inverse=True, return_counts=True), |
83 | | - tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp], npt.NDArray[np.intp]], |
| 82 | + tuple[_nt.Array[np.float64], _nt.Array[np.intp], _nt.Array[np.intp], _nt.Array[np.intp]], |
84 | 83 | ) |
85 | 84 | assert_type( |
86 | 85 | np.unique(AR_LIKE_i_, return_index=True, return_inverse=True, return_counts=True), |
87 | | - tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp], npt.NDArray[np.intp]], |
| 86 | + tuple[_nt.Array[Any], _nt.Array[np.intp], _nt.Array[np.intp], _nt.Array[np.intp]], |
88 | 87 | ) |
89 | 88 |
|
90 | 89 | assert_type(np.unique_all(AR_f8), UniqueAllResult[np.float64]) |
91 | 90 | assert_type(np.unique_counts(AR_f8), UniqueCountsResult[np.float64]) |
92 | 91 | assert_type(np.unique_inverse(AR_f8), UniqueInverseResult[np.float64]) |
93 | | -assert_type(np.unique_values(AR_f8), Array1D[np.float64]) |
| 92 | +assert_type(np.unique_values(AR_f8), _nt.Array1D[np.float64]) |
94 | 93 | assert_type(np.unique_all(AR_LIKE_i_), UniqueAllResult[np.intp]) |
95 | 94 | assert_type(np.unique_counts(AR_LIKE_i_), UniqueCountsResult[np.intp]) |
96 | 95 | assert_type(np.unique_inverse(AR_LIKE_i_), UniqueInverseResult[np.intp]) |
97 | | -assert_type(np.unique_values(AR_LIKE_i_), Array1D[np.intp]) |
| 96 | +assert_type(np.unique_values(AR_LIKE_i_), _nt.Array1D[np.intp]) |
0 commit comments