codekingpro/portable-devtools

1from typing import (
2    Any,
3    Generic,
4    Literal as L,
5    NamedTuple,
6    SupportsIndex,
7    TypeAlias,
8    overload,
9)
10from typing_extensions import TypeVar
11
12import numpy as np
13from numpy._typing import (
14    ArrayLike,
15    NDArray,
16    _ArrayLike,
17    _ArrayLikeBool_co,
18    _ArrayLikeNumber_co,
19)
20
21__all__ = [
22    "ediff1d",
23    "intersect1d",
24    "isin",
25    "setdiff1d",
26    "setxor1d",
27    "union1d",
28    "unique",
29    "unique_all",
30    "unique_counts",
31    "unique_inverse",
32    "unique_values",
33]
34
35_ScalarT = TypeVar("_ScalarT", bound=np.generic)
36_NumericT = TypeVar("_NumericT", bound=np.number | np.timedelta64 | np.object_)
37
38# Explicitly set all allowed values to prevent accidental castings to
39# abstract dtypes (their common super-type).
40# Only relevant if two or more arguments are parametrized, (e.g. `setdiff1d`)
41# which could result in, for example, `int64` and `float64`producing a
42# `number[_64Bit]` array
43_EitherSCT = TypeVar(
44    "_EitherSCT",
45    np.bool,
46    np.int8, np.int16, np.int32, np.int64, np.intp,
47    np.uint8, np.uint16, np.uint32, np.uint64, np.uintp,
48    np.float16, np.float32, np.float64, np.longdouble,
49    np.complex64, np.complex128, np.clongdouble,
50    np.timedelta64, np.datetime64,
51    np.bytes_, np.str_, np.void, np.object_,
52    np.integer, np.floating, np.complexfloating, np.character,
53)  # fmt: skip
54
55_AnyArray: TypeAlias = NDArray[Any]
56_IntArray: TypeAlias = NDArray[np.intp]
57
58###
59
60class UniqueAllResult(NamedTuple, Generic[_ScalarT]):
61    values: NDArray[_ScalarT]
62    indices: _IntArray
63    inverse_indices: _IntArray
64    counts: _IntArray
65
66class UniqueCountsResult(NamedTuple, Generic[_ScalarT]):
67    values: NDArray[_ScalarT]
68    counts: _IntArray
69
70class UniqueInverseResult(NamedTuple, Generic[_ScalarT]):
71    values: NDArray[_ScalarT]
72    inverse_indices: _IntArray
73
74#
75@overload
76def ediff1d(
77    ary: _ArrayLikeBool_co,
78    to_end: ArrayLike | None = None,
79    to_begin: ArrayLike | None = None,
80) -> NDArray[np.int8]: ...
81@overload
82def ediff1d(
83    ary: _ArrayLike[_NumericT],
84    to_end: ArrayLike | None = None,
85    to_begin: ArrayLike | None = None,
86) -> NDArray[_NumericT]: ...
87@overload
88def ediff1d(
89    ary: _ArrayLike[np.datetime64[Any]],
90    to_end: ArrayLike | None = None,
91    to_begin: ArrayLike | None = None,
92) -> NDArray[np.timedelta64]: ...
93@overload
94def ediff1d(
95    ary: _ArrayLikeNumber_co,
96    to_end: ArrayLike | None = None,
97    to_begin: ArrayLike | None = None,
98) -> _AnyArray: ...
99
100#
101@overload  # known scalar-type, FFF
102def unique(
103    ar: _ArrayLike[_ScalarT],
104    return_index: L[False] = False,
105    return_inverse: L[False] = False,
106    return_counts: L[False] = False,
107    axis: SupportsIndex | None = None,
108    *,
109    equal_nan: bool = True,
110    sorted: bool = True,
111) -> NDArray[_ScalarT]: ...
112@overload  # unknown scalar-type, FFF
113def unique(
114    ar: ArrayLike,
115    return_index: L[False] = False,
116    return_inverse: L[False] = False,
117    return_counts: L[False] = False,
118    axis: SupportsIndex | None = None,
119    *,
120    equal_nan: bool = True,
121    sorted: bool = True,
122) -> _AnyArray: ...
123@overload  # known scalar-type, TFF
124def unique(
125    ar: _ArrayLike[_ScalarT],
126    return_index: L[True],
127    return_inverse: L[False] = False,
128    return_counts: L[False] = False,
129    axis: SupportsIndex | None = None,
130    *,
131    equal_nan: bool = True,
132    sorted: bool = True,
133) -> tuple[NDArray[_ScalarT], _IntArray]: ...
134@overload  # unknown scalar-type, TFF
135def unique(
136    ar: ArrayLike,
137    return_index: L[True],
138    return_inverse: L[False] = False,
139    return_counts: L[False] = False,
140    axis: SupportsIndex | None = None,
141    *,
142    equal_nan: bool = True,
143    sorted: bool = True,
144) -> tuple[_AnyArray, _IntArray]: ...
145@overload  # known scalar-type, FTF (positional)
146def unique(
147    ar: _ArrayLike[_ScalarT],
148    return_index: L[False],
149    return_inverse: L[True],
150    return_counts: L[False] = False,
151    axis: SupportsIndex | None = None,
152    *,
153    equal_nan: bool = True,
154    sorted: bool = True,
155) -> tuple[NDArray[_ScalarT], _IntArray]: ...
156@overload  # known scalar-type, FTF (keyword)
157def unique(
158    ar: _ArrayLike[_ScalarT],
159    return_index: L[False] = False,
160    *,
161    return_inverse: L[True],
162    return_counts: L[False] = False,
163    axis: SupportsIndex | None = None,
164    equal_nan: bool = True,
165    sorted: bool = True,
166) -> tuple[NDArray[_ScalarT], _IntArray]: ...
167@overload  # unknown scalar-type, FTF (positional)
168def unique(
169    ar: ArrayLike,
170    return_index: L[False],
171    return_inverse: L[True],
172    return_counts: L[False] = False,
173    axis: SupportsIndex | None = None,
174    *,
175    equal_nan: bool = True,
176    sorted: bool = True,
177) -> tuple[_AnyArray, _IntArray]: ...
178@overload  # unknown scalar-type, FTF (keyword)
179def unique(
180    ar: ArrayLike,
181    return_index: L[False] = False,
182    *,
183    return_inverse: L[True],
184    return_counts: L[False] = False,
185    axis: SupportsIndex | None = None,
186    equal_nan: bool = True,
187    sorted: bool = True,
188) -> tuple[_AnyArray, _IntArray]: ...
189@overload  # known scalar-type, FFT (positional)
190def unique(
191    ar: _ArrayLike[_ScalarT],
192    return_index: L[False],
193    return_inverse: L[False],
194    return_counts: L[True],
195    axis: SupportsIndex | None = None,
196    *,
197    equal_nan: bool = True,
198    sorted: bool = True,
199) -> tuple[NDArray[_ScalarT], _IntArray]: ...
200@overload  # known scalar-type, FFT (keyword)
201def unique(
202    ar: _ArrayLike[_ScalarT],
203    return_index: L[False] = False,
204    return_inverse: L[False] = False,
205    *,
206    return_counts: L[True],
207    axis: SupportsIndex | None = None,
208    equal_nan: bool = True,
209    sorted: bool = True,
210) -> tuple[NDArray[_ScalarT], _IntArray]: ...
211@overload  # unknown scalar-type, FFT (positional)
212def unique(
213    ar: ArrayLike,
214    return_index: L[False],
215    return_inverse: L[False],
216    return_counts: L[True],
217    axis: SupportsIndex | None = None,
218    *,
219    equal_nan: bool = True,
220    sorted: bool = True,
221) -> tuple[_AnyArray, _IntArray]: ...
222@overload  # unknown scalar-type, FFT (keyword)
223def unique(
224    ar: ArrayLike,
225    return_index: L[False] = False,
226    return_inverse: L[False] = False,
227    *,
228    return_counts: L[True],
229    axis: SupportsIndex | None = None,
230    equal_nan: bool = True,
231    sorted: bool = True,
232) -> tuple[_AnyArray, _IntArray]: ...
233@overload  # known scalar-type, TTF
234def unique(
235    ar: _ArrayLike[_ScalarT],
236    return_index: L[True],
237    return_inverse: L[True],
238    return_counts: L[False] = False,
239    axis: SupportsIndex | None = None,
240    *,
241    equal_nan: bool = True,
242    sorted: bool = True,
243) -> tuple[NDArray[_ScalarT], _IntArray, _IntArray]: ...
244@overload  # unknown scalar-type, TTF
245def unique(
246    ar: ArrayLike,
247    return_index: L[True],
248    return_inverse: L[True],
249    return_counts: L[False] = False,
250    axis: SupportsIndex | None = None,
251    *,
252    equal_nan: bool = True,
253    sorted: bool = True,
254) -> tuple[_AnyArray, _IntArray, _IntArray]: ...
255@overload  # known scalar-type, TFT (positional)
256def unique(
257    ar: _ArrayLike[_ScalarT],
258    return_index: L[True],
259    return_inverse: L[False],
260    return_counts: L[True],
261    axis: SupportsIndex | None = None,
262    *,
263    equal_nan: bool = True,
264    sorted: bool = True,
265) -> tuple[NDArray[_ScalarT], _IntArray, _IntArray]: ...
266@overload  # known scalar-type, TFT (keyword)
267def unique(
268    ar: _ArrayLike[_ScalarT],
269    return_index: L[True],
270    return_inverse: L[False] = False,
271    *,
272    return_counts: L[True],
273    axis: SupportsIndex | None = None,
274    equal_nan: bool = True,
275    sorted: bool = True,
276) -> tuple[NDArray[_ScalarT], _IntArray, _IntArray]: ...
277@overload  # unknown scalar-type, TFT (positional)
278def unique(
279    ar: ArrayLike,
280    return_index: L[True],
281    return_inverse: L[False],
282    return_counts: L[True],
283    axis: SupportsIndex | None = None,
284    *,
285    equal_nan: bool = True,
286    sorted: bool = True,
287) -> tuple[_AnyArray, _IntArray, _IntArray]: ...
288@overload  # unknown scalar-type, TFT (keyword)
289def unique(
290    ar: ArrayLike,
291    return_index: L[True],
292    return_inverse: L[False] = False,
293    *,
294    return_counts: L[True],
295    axis: SupportsIndex | None = None,
296    equal_nan: bool = True,
297    sorted: bool = True,
298) -> tuple[_AnyArray, _IntArray, _IntArray]: ...
299@overload  # known scalar-type, FTT (positional)
300def unique(
301    ar: _ArrayLike[_ScalarT],
302    return_index: L[False],
303    return_inverse: L[True],
304    return_counts: L[True],
305    axis: SupportsIndex | None = None,
306    *,
307    equal_nan: bool = True,
308    sorted: bool = True,
309) -> tuple[NDArray[_ScalarT], _IntArray, _IntArray]: ...
310@overload  # known scalar-type, FTT (keyword)
311def unique(
312    ar: _ArrayLike[_ScalarT],
313    return_index: L[False] = False,
314    *,
315    return_inverse: L[True],
316    return_counts: L[True],
317    axis: SupportsIndex | None = None,
318    equal_nan: bool = True,
319    sorted: bool = True,
320) -> tuple[NDArray[_ScalarT], _IntArray, _IntArray]: ...
321@overload  # unknown scalar-type, FTT (positional)
322def unique(
323    ar: ArrayLike,
324    return_index: L[False],
325    return_inverse: L[True],
326    return_counts: L[True],
327    axis: SupportsIndex | None = None,
328    *,
329    equal_nan: bool = True,
330    sorted: bool = True,
331) -> tuple[_AnyArray, _IntArray, _IntArray]: ...
332@overload  # unknown scalar-type, FTT (keyword)
333def unique(
334    ar: ArrayLike,
335    return_index: L[False] = False,
336    *,
337    return_inverse: L[True],
338    return_counts: L[True],
339    axis: SupportsIndex | None = None,
340    equal_nan: bool = True,
341    sorted: bool = True,
342) -> tuple[_AnyArray, _IntArray, _IntArray]: ...
343@overload  # known scalar-type, TTT
344def unique(
345    ar: _ArrayLike[_ScalarT],
346    return_index: L[True],
347    return_inverse: L[True],
348    return_counts: L[True],
349    axis: SupportsIndex | None = None,
350    *,
351    equal_nan: bool = True,
352    sorted: bool = True,
353) -> tuple[NDArray[_ScalarT], _IntArray, _IntArray, _IntArray]: ...
354@overload  # unknown scalar-type, TTT
355def unique(
356    ar: ArrayLike,
357    return_index: L[True],
358    return_inverse: L[True],
359    return_counts: L[True],
360    axis: SupportsIndex | None = None,
361    *,
362    equal_nan: bool = True,
363    sorted: bool = True,
364) -> tuple[_AnyArray, _IntArray, _IntArray, _IntArray]: ...
365
366#
367@overload
368def unique_all(x: _ArrayLike[_ScalarT]) -> UniqueAllResult[_ScalarT]: ...
369@overload
370def unique_all(x: ArrayLike) -> UniqueAllResult[Any]: ...
371
372#
373@overload
374def unique_counts(x: _ArrayLike[_ScalarT]) -> UniqueCountsResult[_ScalarT]: ...
375@overload
376def unique_counts(x: ArrayLike) -> UniqueCountsResult[Any]: ...
377
378#
379@overload
380def unique_inverse(x: _ArrayLike[_ScalarT]) -> UniqueInverseResult[_ScalarT]: ...
381@overload
382def unique_inverse(x: ArrayLike) -> UniqueInverseResult[Any]: ...
383
384#
385@overload
386def unique_values(x: _ArrayLike[_ScalarT]) -> NDArray[_ScalarT]: ...
387@overload
388def unique_values(x: ArrayLike) -> _AnyArray: ...
389
390#
391@overload  # known scalar-type, return_indices=False (default)
392def intersect1d(
393    ar1: _ArrayLike[_EitherSCT],
394    ar2: _ArrayLike[_EitherSCT],
395    assume_unique: bool = False,
396    return_indices: L[False] = False,
397) -> NDArray[_EitherSCT]: ...
398@overload  # known scalar-type, return_indices=True (positional)
399def intersect1d(
400    ar1: _ArrayLike[_EitherSCT],
401    ar2: _ArrayLike[_EitherSCT],
402    assume_unique: bool,
403    return_indices: L[True],
404) -> tuple[NDArray[_EitherSCT], _IntArray, _IntArray]: ...
405@overload  # known scalar-type, return_indices=True (keyword)
406def intersect1d(
407    ar1: _ArrayLike[_EitherSCT],
408    ar2: _ArrayLike[_EitherSCT],
409    assume_unique: bool = False,
410    *,
411    return_indices: L[True],
412) -> tuple[NDArray[_EitherSCT], _IntArray, _IntArray]: ...
413@overload  # unknown scalar-type, return_indices=False (default)
414def intersect1d(
415    ar1: ArrayLike,
416    ar2: ArrayLike,
417    assume_unique: bool = False,
418    return_indices: L[False] = False,
419) -> _AnyArray: ...
420@overload  # unknown scalar-type, return_indices=True (positional)
421def intersect1d(
422    ar1: ArrayLike,
423    ar2: ArrayLike,
424    assume_unique: bool,
425    return_indices: L[True],
426) -> tuple[_AnyArray, _IntArray, _IntArray]: ...
427@overload  # unknown scalar-type, return_indices=True (keyword)
428def intersect1d(
429    ar1: ArrayLike,
430    ar2: ArrayLike,
431    assume_unique: bool = False,
432    *,
433    return_indices: L[True],
434) -> tuple[_AnyArray, _IntArray, _IntArray]: ...
435
436#
437@overload
438def setxor1d(ar1: _ArrayLike[_EitherSCT], ar2: _ArrayLike[_EitherSCT], assume_unique: bool = False) -> NDArray[_EitherSCT]: ...
439@overload
440def setxor1d(ar1: ArrayLike, ar2: ArrayLike, assume_unique: bool = False) -> _AnyArray: ...
441
442#
443@overload
444def union1d(ar1: _ArrayLike[_EitherSCT], ar2: _ArrayLike[_EitherSCT]) -> NDArray[_EitherSCT]: ...
445@overload
446def union1d(ar1: ArrayLike, ar2: ArrayLike) -> _AnyArray: ...
447
448#
449@overload
450def setdiff1d(ar1: _ArrayLike[_EitherSCT], ar2: _ArrayLike[_EitherSCT], assume_unique: bool = False) -> NDArray[_EitherSCT]: ...
451@overload
452def setdiff1d(ar1: ArrayLike, ar2: ArrayLike, assume_unique: bool = False) -> _AnyArray: ...
453
454#
455def isin(
456    element: ArrayLike,
457    test_elements: ArrayLike,
458    assume_unique: bool = False,
459    invert: bool = False,
460    *,
461    kind: L["sort", "table"] | None = None,
462) -> NDArray[np.bool]: ...
463