_arraysetops_impl.pyi 13 KB

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