Skip to content

Commit 5caaabf

Browse files
author
Bas van Beek
committed
ENH: Add annotations for np.lib.arraysetops 
1 parent e467a28 commit 5caaabf

File tree

1 file changed

+332
-9
lines changed

1 file changed

+332
-9
lines changed

‎numpy/lib/arraysetops.pyi‎

Lines changed: 332 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,335 @@
1-
from typing import List
1+
from typing import (
2+
Literal as L,
3+
Any,
4+
List,
5+
Union,
6+
TypeVar,
7+
Tuple,
8+
overload,
9+
SupportsIndex,
10+
)
11+
12+
from numpy import (
13+
dtype,
14+
generic,
15+
number,
16+
bool_,
17+
ushort,
18+
ubyte,
19+
uintc,
20+
uint,
21+
ulonglong,
22+
short,
23+
int8,
24+
byte,
25+
intc,
26+
int_,
27+
intp,
28+
longlong,
29+
half,
30+
single,
31+
double,
32+
longdouble,
33+
csingle,
34+
cdouble,
35+
clongdouble,
36+
timedelta64,
37+
datetime64,
38+
object_,
39+
str_,
40+
bytes_,
41+
void,
42+
)
43+
44+
from numpy.typing import (
45+
ArrayLike,
46+
NDArray,
47+
_FiniteNestedSequence,
48+
_SupportsArray,
49+
_ArrayLikeBool_co,
50+
_ArrayLikeDT64_co,
51+
_ArrayLikeTD64_co,
52+
_ArrayLikeObject_co,
53+
_ArrayLikeNumber_co,
54+
)
55+
56+
_SCT = TypeVar("_SCT", bound=generic)
57+
_NumberType = TypeVar("_NumberType", bound=number[Any])
58+
59+
# Explicitly set all allowed values to prevent accidental castings to
60+
# abstract dtypes (their common super-type).
61+
#
62+
# Only relevant if two or more arguments are parametrized, (e.g. `setdiff1d`)
63+
# which could result in, for example, `int64` and `float64`producing a
64+
# `number[_64Bit]` array
65+
_SCTNoCast = TypeVar(
66+
"_SCTNoCast",
67+
bool_,
68+
ushort,
69+
ubyte,
70+
uintc,
71+
uint,
72+
ulonglong,
73+
short,
74+
byte,
75+
intc,
76+
int_,
77+
longlong,
78+
half,
79+
single,
80+
double,
81+
longdouble,
82+
csingle,
83+
cdouble,
84+
clongdouble,
85+
timedelta64,
86+
datetime64,
87+
object_,
88+
str_,
89+
bytes_,
90+
void,
91+
)
92+
93+
_ArrayLike = _FiniteNestedSequence[_SupportsArray[dtype[_SCT]]]
294

395
__all__: List[str]
496

5-
def ediff1d(ary, to_end=..., to_begin=...): ...
6-
def unique(ar, return_index=..., return_inverse=..., return_counts=..., axis=...): ...
7-
def intersect1d(ar1, ar2, assume_unique=..., return_indices=...): ...
8-
def setxor1d(ar1, ar2, assume_unique=...): ...
9-
def in1d(ar1, ar2, assume_unique=..., invert=...): ...
10-
def isin(element, test_elements, assume_unique=..., invert=...): ...
11-
def union1d(ar1, ar2): ...
12-
def setdiff1d(ar1, ar2, assume_unique=...): ...
97+
@overload
98+
def ediff1d(
99+
ary: _ArrayLikeBool_co,
100+
to_end: None | ArrayLike = ...,
101+
to_begin: None | ArrayLike = ...,
102+
) -> NDArray[int8]: ...
103+
@overload
104+
def ediff1d(
105+
ary: _ArrayLike[_NumberType],
106+
to_end: None | ArrayLike = ...,
107+
to_begin: None | ArrayLike = ...,
108+
) -> NDArray[_NumberType]: ...
109+
@overload
110+
def ediff1d(
111+
ary: _ArrayLikeNumber_co,
112+
to_end: None | ArrayLike = ...,
113+
to_begin: None | ArrayLike = ...,
114+
) -> NDArray[Any]: ...
115+
@overload
116+
def ediff1d(
117+
ary: _ArrayLikeDT64_co | _ArrayLikeTD64_co,
118+
to_end: None | ArrayLike = ...,
119+
to_begin: None | ArrayLike = ...,
120+
) -> NDArray[timedelta64]: ...
121+
@overload
122+
def ediff1d(
123+
ary: _ArrayLikeObject_co,
124+
to_end: None | ArrayLike = ...,
125+
to_begin: None | ArrayLike = ...,
126+
) -> NDArray[object_]: ...
127+
128+
@overload
129+
def unique(
130+
ar: _ArrayLike[_SCT],
131+
return_index: L[False] = ...,
132+
return_inverse: L[False] = ...,
133+
return_counts: L[False] = ...,
134+
axis: None | SupportsIndex = ...,
135+
) -> NDArray[_SCT]: ...
136+
@overload
137+
def unique(
138+
ar: ArrayLike,
139+
return_index: L[False] = ...,
140+
return_inverse: L[False] = ...,
141+
return_counts: L[False] = ...,
142+
axis: None | SupportsIndex = ...,
143+
) -> NDArray[Any]: ...
144+
@overload
145+
def unique(
146+
ar: _ArrayLike[_SCT],
147+
return_index: L[True] = ...,
148+
return_inverse: L[False] = ...,
149+
return_counts: L[False] = ...,
150+
axis: None | SupportsIndex = ...,
151+
) -> Tuple[NDArray[_SCT], NDArray[intp]]: ...
152+
@overload
153+
def unique(
154+
ar: ArrayLike,
155+
return_index: L[True] = ...,
156+
return_inverse: L[False] = ...,
157+
return_counts: L[False] = ...,
158+
axis: None | SupportsIndex = ...,
159+
) -> Tuple[NDArray[Any], NDArray[intp]]: ...
160+
@overload
161+
def unique(
162+
ar: _ArrayLike[_SCT],
163+
return_index: L[False] = ...,
164+
return_inverse: L[True] = ...,
165+
return_counts: L[False] = ...,
166+
axis: None | SupportsIndex = ...,
167+
) -> Tuple[NDArray[_SCT], NDArray[intp]]: ...
168+
@overload
169+
def unique(
170+
ar: ArrayLike,
171+
return_index: L[False] = ...,
172+
return_inverse: L[True] = ...,
173+
return_counts: L[False] = ...,
174+
axis: None | SupportsIndex = ...,
175+
) -> Tuple[NDArray[Any], NDArray[intp]]: ...
176+
@overload
177+
def unique(
178+
ar: _ArrayLike[_SCT],
179+
return_index: L[False] = ...,
180+
return_inverse: L[False] = ...,
181+
return_counts: L[True] = ...,
182+
axis: None | SupportsIndex = ...,
183+
) -> Tuple[NDArray[_SCT], NDArray[intp]]: ...
184+
@overload
185+
def unique(
186+
ar: ArrayLike,
187+
return_index: L[False] = ...,
188+
return_inverse: L[False] = ...,
189+
return_counts: L[True] = ...,
190+
axis: None | SupportsIndex = ...,
191+
) -> Tuple[NDArray[Any], NDArray[intp]]: ...
192+
@overload
193+
def unique(
194+
ar: _ArrayLike[_SCT],
195+
return_index: L[True] = ...,
196+
return_inverse: L[True] = ...,
197+
return_counts: L[False] = ...,
198+
axis: None | SupportsIndex = ...,
199+
) -> Tuple[NDArray[_SCT], NDArray[intp], NDArray[intp]]: ...
200+
@overload
201+
def unique(
202+
ar: ArrayLike,
203+
return_index: L[True] = ...,
204+
return_inverse: L[True] = ...,
205+
return_counts: L[False] = ...,
206+
axis: None | SupportsIndex = ...,
207+
) -> Tuple[NDArray[Any], NDArray[intp], NDArray[intp]]: ...
208+
@overload
209+
def unique(
210+
ar: _ArrayLike[_SCT],
211+
return_index: L[True] = ...,
212+
return_inverse: L[False] = ...,
213+
return_counts: L[True] = ...,
214+
axis: None | SupportsIndex = ...,
215+
) -> Tuple[NDArray[_SCT], NDArray[intp], NDArray[intp]]: ...
216+
@overload
217+
def unique(
218+
ar: ArrayLike,
219+
return_index: L[True] = ...,
220+
return_inverse: L[False] = ...,
221+
return_counts: L[True] = ...,
222+
axis: None | SupportsIndex = ...,
223+
) -> Tuple[NDArray[Any], NDArray[intp], NDArray[intp]]: ...
224+
@overload
225+
def unique(
226+
ar: _ArrayLike[_SCT],
227+
return_index: L[False] = ...,
228+
return_inverse: L[True] = ...,
229+
return_counts: L[True] = ...,
230+
axis: None | SupportsIndex = ...,
231+
) -> Tuple[NDArray[_SCT], NDArray[intp], NDArray[intp]]: ...
232+
@overload
233+
def unique(
234+
ar: ArrayLike,
235+
return_index: L[False] = ...,
236+
return_inverse: L[True] = ...,
237+
return_counts: L[True] = ...,
238+
axis: None | SupportsIndex = ...,
239+
) -> Tuple[NDArray[Any], NDArray[intp], NDArray[intp]]: ...
240+
@overload
241+
def unique(
242+
ar: _ArrayLike[_SCT],
243+
return_index: L[True] = ...,
244+
return_inverse: L[True] = ...,
245+
return_counts: L[True] = ...,
246+
axis: None | SupportsIndex = ...,
247+
) -> Tuple[NDArray[_SCT], NDArray[intp], NDArray[intp], NDArray[intp]]: ...
248+
@overload
249+
def unique(
250+
ar: ArrayLike,
251+
return_index: L[True] = ...,
252+
return_inverse: L[True] = ...,
253+
return_counts: L[True] = ...,
254+
axis: None | SupportsIndex = ...,
255+
) -> Tuple[NDArray[Any], NDArray[intp], NDArray[intp], NDArray[intp]]: ...
256+
257+
@overload
258+
def intersect1d(
259+
ar1: _ArrayLike[_SCTNoCast],
260+
ar2: _ArrayLike[_SCTNoCast],
261+
assume_unique: bool = ...,
262+
return_indices: L[False] = ...,
263+
) -> NDArray[_SCTNoCast]: ...
264+
@overload
265+
def intersect1d(
266+
ar1: ArrayLike,
267+
ar2: ArrayLike,
268+
assume_unique: bool = ...,
269+
return_indices: L[False] = ...,
270+
) -> NDArray[Any]: ...
271+
@overload
272+
def intersect1d(
273+
ar1: _ArrayLike[_SCTNoCast],
274+
ar2: _ArrayLike[_SCTNoCast],
275+
assume_unique: bool = ...,
276+
return_indices: L[True] = ...,
277+
) -> Tuple[NDArray[_SCTNoCast], NDArray[intp], NDArray[intp]]: ...
278+
@overload
279+
def intersect1d(
280+
ar1: ArrayLike,
281+
ar2: ArrayLike,
282+
assume_unique: bool = ...,
283+
return_indices: L[True] = ...,
284+
) -> Tuple[NDArray[Any], NDArray[intp], NDArray[intp]]: ...
285+
286+
@overload
287+
def setxor1d(
288+
ar1: _ArrayLike[_SCTNoCast],
289+
ar2: _ArrayLike[_SCTNoCast],
290+
assume_unique: bool = ...,
291+
) -> NDArray[_SCTNoCast]: ...
292+
@overload
293+
def setxor1d(
294+
ar1: ArrayLike,
295+
ar2: ArrayLike,
296+
assume_unique: bool = ...,
297+
) -> NDArray[Any]: ...
298+
299+
def in1d(
300+
ar1: ArrayLike,
301+
ar2: ArrayLike,
302+
assume_unique: bool = ...,
303+
invert: bool = ...,
304+
) -> NDArray[bool_]: ...
305+
306+
def isin(
307+
element: ArrayLike,
308+
test_elements: ArrayLike,
309+
assume_unique: bool = ...,
310+
invert: bool = ...,
311+
) -> NDArray[bool_]: ...
312+
313+
@overload
314+
def union1d(
315+
ar1: _ArrayLike[_SCTNoCast],
316+
ar2: _ArrayLike[_SCTNoCast],
317+
) -> NDArray[_SCTNoCast]: ...
318+
@overload
319+
def union1d(
320+
ar1: ArrayLike,
321+
ar2: ArrayLike,
322+
) -> NDArray[Any]: ...
323+
324+
@overload
325+
def setdiff1d(
326+
ar1: _ArrayLike[_SCTNoCast],
327+
ar2: _ArrayLike[_SCTNoCast],
328+
assume_unique: bool = ...,
329+
) -> NDArray[_SCTNoCast]: ...
330+
@overload
331+
def setdiff1d(
332+
ar1: ArrayLike,
333+
ar2: ArrayLike,
334+
assume_unique: bool = ...,
335+
) -> NDArray[Any]: ...

0 commit comments

Comments
 (0)