_arraysetops_impl.pyi 9.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379
  1. from typing import (
  2. Any,
  3. Generic,
  4. Literal as L,
  5. NamedTuple,
  6. overload,
  7. SupportsIndex,
  8. TypeVar,
  9. )
  10. from typing_extensions import deprecated
  11. import numpy as np
  12. from numpy import generic, number, int8, intp, timedelta64, object_
  13. from numpy._typing import (
  14. ArrayLike,
  15. NDArray,
  16. _ArrayLike,
  17. _ArrayLikeBool_co,
  18. _ArrayLikeDT64_co,
  19. _ArrayLikeTD64_co,
  20. _ArrayLikeObject_co,
  21. _ArrayLikeNumber_co,
  22. )
  23. __all__ = [
  24. "ediff1d",
  25. "in1d",
  26. "intersect1d",
  27. "isin",
  28. "setdiff1d",
  29. "setxor1d",
  30. "union1d",
  31. "unique",
  32. "unique_all",
  33. "unique_counts",
  34. "unique_inverse",
  35. "unique_values",
  36. ]
  37. _SCT = TypeVar("_SCT", bound=generic)
  38. _NumberType = TypeVar("_NumberType", bound=number[Any])
  39. # Explicitly set all allowed values to prevent accidental castings to
  40. # abstract dtypes (their common super-type).
  41. #
  42. # Only relevant if two or more arguments are parametrized, (e.g. `setdiff1d`)
  43. # which could result in, for example, `int64` and `float64`producing a
  44. # `number[_64Bit]` array
  45. _EitherSCT = TypeVar(
  46. "_EitherSCT",
  47. np.bool,
  48. np.int8, np.int16, np.int32, np.int64, np.intp,
  49. np.uint8, np.uint16, np.uint32, np.uint64, np.uintp,
  50. np.float16, np.float32, np.float64, np.longdouble,
  51. np.complex64, np.complex128, np.clongdouble,
  52. np.timedelta64, np.datetime64,
  53. np.bytes_, np.str_, np.void, np.object_,
  54. np.integer, np.floating, np.complexfloating, np.character,
  55. ) # fmt: skip
  56. class UniqueAllResult(NamedTuple, Generic[_SCT]):
  57. values: NDArray[_SCT]
  58. indices: NDArray[intp]
  59. inverse_indices: NDArray[intp]
  60. counts: NDArray[intp]
  61. class UniqueCountsResult(NamedTuple, Generic[_SCT]):
  62. values: NDArray[_SCT]
  63. counts: NDArray[intp]
  64. class UniqueInverseResult(NamedTuple, Generic[_SCT]):
  65. values: NDArray[_SCT]
  66. inverse_indices: NDArray[intp]
  67. @overload
  68. def ediff1d(
  69. ary: _ArrayLikeBool_co,
  70. to_end: None | ArrayLike = ...,
  71. to_begin: None | ArrayLike = ...,
  72. ) -> NDArray[int8]: ...
  73. @overload
  74. def ediff1d(
  75. ary: _ArrayLike[_NumberType],
  76. to_end: None | ArrayLike = ...,
  77. to_begin: None | ArrayLike = ...,
  78. ) -> NDArray[_NumberType]: ...
  79. @overload
  80. def ediff1d(
  81. ary: _ArrayLikeNumber_co,
  82. to_end: None | ArrayLike = ...,
  83. to_begin: None | ArrayLike = ...,
  84. ) -> NDArray[Any]: ...
  85. @overload
  86. def ediff1d(
  87. ary: _ArrayLikeDT64_co | _ArrayLikeTD64_co,
  88. to_end: None | ArrayLike = ...,
  89. to_begin: None | ArrayLike = ...,
  90. ) -> NDArray[timedelta64]: ...
  91. @overload
  92. def ediff1d(
  93. ary: _ArrayLikeObject_co,
  94. to_end: None | ArrayLike = ...,
  95. to_begin: None | ArrayLike = ...,
  96. ) -> NDArray[object_]: ...
  97. @overload
  98. def unique(
  99. ar: _ArrayLike[_SCT],
  100. return_index: L[False] = ...,
  101. return_inverse: L[False] = ...,
  102. return_counts: L[False] = ...,
  103. axis: None | SupportsIndex = ...,
  104. *,
  105. equal_nan: bool = ...,
  106. ) -> NDArray[_SCT]: ...
  107. @overload
  108. def unique(
  109. ar: ArrayLike,
  110. return_index: L[False] = ...,
  111. return_inverse: L[False] = ...,
  112. return_counts: L[False] = ...,
  113. axis: None | SupportsIndex = ...,
  114. *,
  115. equal_nan: bool = ...,
  116. ) -> NDArray[Any]: ...
  117. @overload
  118. def unique(
  119. ar: _ArrayLike[_SCT],
  120. return_index: L[True] = ...,
  121. return_inverse: L[False] = ...,
  122. return_counts: L[False] = ...,
  123. axis: None | SupportsIndex = ...,
  124. *,
  125. equal_nan: bool = ...,
  126. ) -> tuple[NDArray[_SCT], NDArray[intp]]: ...
  127. @overload
  128. def unique(
  129. ar: ArrayLike,
  130. return_index: L[True] = ...,
  131. return_inverse: L[False] = ...,
  132. return_counts: L[False] = ...,
  133. axis: None | SupportsIndex = ...,
  134. *,
  135. equal_nan: bool = ...,
  136. ) -> tuple[NDArray[Any], NDArray[intp]]: ...
  137. @overload
  138. def unique(
  139. ar: _ArrayLike[_SCT],
  140. return_index: L[False] = ...,
  141. return_inverse: L[True] = ...,
  142. return_counts: L[False] = ...,
  143. axis: None | SupportsIndex = ...,
  144. *,
  145. equal_nan: bool = ...,
  146. ) -> tuple[NDArray[_SCT], NDArray[intp]]: ...
  147. @overload
  148. def unique(
  149. ar: ArrayLike,
  150. return_index: L[False] = ...,
  151. return_inverse: L[True] = ...,
  152. return_counts: L[False] = ...,
  153. axis: None | SupportsIndex = ...,
  154. *,
  155. equal_nan: bool = ...,
  156. ) -> tuple[NDArray[Any], NDArray[intp]]: ...
  157. @overload
  158. def unique(
  159. ar: _ArrayLike[_SCT],
  160. return_index: L[False] = ...,
  161. return_inverse: L[False] = ...,
  162. return_counts: L[True] = ...,
  163. axis: None | SupportsIndex = ...,
  164. *,
  165. equal_nan: bool = ...,
  166. ) -> tuple[NDArray[_SCT], NDArray[intp]]: ...
  167. @overload
  168. def unique(
  169. ar: ArrayLike,
  170. return_index: L[False] = ...,
  171. return_inverse: L[False] = ...,
  172. return_counts: L[True] = ...,
  173. axis: None | SupportsIndex = ...,
  174. *,
  175. equal_nan: bool = ...,
  176. ) -> tuple[NDArray[Any], NDArray[intp]]: ...
  177. @overload
  178. def unique(
  179. ar: _ArrayLike[_SCT],
  180. return_index: L[True] = ...,
  181. return_inverse: L[True] = ...,
  182. return_counts: L[False] = ...,
  183. axis: None | SupportsIndex = ...,
  184. *,
  185. equal_nan: bool = ...,
  186. ) -> tuple[NDArray[_SCT], NDArray[intp], NDArray[intp]]: ...
  187. @overload
  188. def unique(
  189. ar: ArrayLike,
  190. return_index: L[True] = ...,
  191. return_inverse: L[True] = ...,
  192. return_counts: L[False] = ...,
  193. axis: None | SupportsIndex = ...,
  194. *,
  195. equal_nan: bool = ...,
  196. ) -> tuple[NDArray[Any], NDArray[intp], NDArray[intp]]: ...
  197. @overload
  198. def unique(
  199. ar: _ArrayLike[_SCT],
  200. return_index: L[True] = ...,
  201. return_inverse: L[False] = ...,
  202. return_counts: L[True] = ...,
  203. axis: None | SupportsIndex = ...,
  204. *,
  205. equal_nan: bool = ...,
  206. ) -> tuple[NDArray[_SCT], NDArray[intp], NDArray[intp]]: ...
  207. @overload
  208. def unique(
  209. ar: ArrayLike,
  210. return_index: L[True] = ...,
  211. return_inverse: L[False] = ...,
  212. return_counts: L[True] = ...,
  213. axis: None | SupportsIndex = ...,
  214. *,
  215. equal_nan: bool = ...,
  216. ) -> tuple[NDArray[Any], NDArray[intp], NDArray[intp]]: ...
  217. @overload
  218. def unique(
  219. ar: _ArrayLike[_SCT],
  220. return_index: L[False] = ...,
  221. return_inverse: L[True] = ...,
  222. return_counts: L[True] = ...,
  223. axis: None | SupportsIndex = ...,
  224. *,
  225. equal_nan: bool = ...,
  226. ) -> tuple[NDArray[_SCT], NDArray[intp], NDArray[intp]]: ...
  227. @overload
  228. def unique(
  229. ar: ArrayLike,
  230. return_index: L[False] = ...,
  231. return_inverse: L[True] = ...,
  232. return_counts: L[True] = ...,
  233. axis: None | SupportsIndex = ...,
  234. *,
  235. equal_nan: bool = ...,
  236. ) -> tuple[NDArray[Any], NDArray[intp], NDArray[intp]]: ...
  237. @overload
  238. def unique(
  239. ar: _ArrayLike[_SCT],
  240. return_index: L[True] = ...,
  241. return_inverse: L[True] = ...,
  242. return_counts: L[True] = ...,
  243. axis: None | SupportsIndex = ...,
  244. *,
  245. equal_nan: bool = ...,
  246. ) -> tuple[NDArray[_SCT], NDArray[intp], NDArray[intp], NDArray[intp]]: ...
  247. @overload
  248. def unique(
  249. ar: ArrayLike,
  250. return_index: L[True] = ...,
  251. return_inverse: L[True] = ...,
  252. return_counts: L[True] = ...,
  253. axis: None | SupportsIndex = ...,
  254. *,
  255. equal_nan: bool = ...,
  256. ) -> tuple[NDArray[Any], NDArray[intp], NDArray[intp], NDArray[intp]]: ...
  257. @overload
  258. def unique_all(
  259. x: _ArrayLike[_SCT], /
  260. ) -> UniqueAllResult[_SCT]: ...
  261. @overload
  262. def unique_all(
  263. x: ArrayLike, /
  264. ) -> UniqueAllResult[Any]: ...
  265. @overload
  266. def unique_counts(
  267. x: _ArrayLike[_SCT], /
  268. ) -> UniqueCountsResult[_SCT]: ...
  269. @overload
  270. def unique_counts(
  271. x: ArrayLike, /
  272. ) -> UniqueCountsResult[Any]: ...
  273. @overload
  274. def unique_inverse(x: _ArrayLike[_SCT], /) -> UniqueInverseResult[_SCT]: ...
  275. @overload
  276. def unique_inverse(x: ArrayLike, /) -> UniqueInverseResult[Any]: ...
  277. @overload
  278. def unique_values(x: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
  279. @overload
  280. def unique_values(x: ArrayLike, /) -> NDArray[Any]: ...
  281. @overload
  282. def intersect1d(
  283. ar1: _ArrayLike[_EitherSCT],
  284. ar2: _ArrayLike[_EitherSCT],
  285. assume_unique: bool = ...,
  286. return_indices: L[False] = ...,
  287. ) -> NDArray[_EitherSCT]: ...
  288. @overload
  289. def intersect1d(
  290. ar1: ArrayLike,
  291. ar2: ArrayLike,
  292. assume_unique: bool = ...,
  293. return_indices: L[False] = ...,
  294. ) -> NDArray[Any]: ...
  295. @overload
  296. def intersect1d(
  297. ar1: _ArrayLike[_EitherSCT],
  298. ar2: _ArrayLike[_EitherSCT],
  299. assume_unique: bool = ...,
  300. return_indices: L[True] = ...,
  301. ) -> tuple[NDArray[_EitherSCT], NDArray[intp], NDArray[intp]]: ...
  302. @overload
  303. def intersect1d(
  304. ar1: ArrayLike,
  305. ar2: ArrayLike,
  306. assume_unique: bool = ...,
  307. return_indices: L[True] = ...,
  308. ) -> tuple[NDArray[Any], NDArray[intp], NDArray[intp]]: ...
  309. @overload
  310. def setxor1d(
  311. ar1: _ArrayLike[_EitherSCT],
  312. ar2: _ArrayLike[_EitherSCT],
  313. assume_unique: bool = ...,
  314. ) -> NDArray[_EitherSCT]: ...
  315. @overload
  316. def setxor1d(
  317. ar1: ArrayLike,
  318. ar2: ArrayLike,
  319. assume_unique: bool = ...,
  320. ) -> NDArray[Any]: ...
  321. def isin(
  322. element: ArrayLike,
  323. test_elements: ArrayLike,
  324. assume_unique: bool = ...,
  325. invert: bool = ...,
  326. *,
  327. kind: None | str = ...,
  328. ) -> NDArray[np.bool]: ...
  329. @deprecated("Use 'isin' instead")
  330. def in1d(
  331. element: ArrayLike,
  332. test_elements: ArrayLike,
  333. assume_unique: bool = ...,
  334. invert: bool = ...,
  335. *,
  336. kind: None | str = ...,
  337. ) -> NDArray[np.bool]: ...
  338. @overload
  339. def union1d(
  340. ar1: _ArrayLike[_EitherSCT],
  341. ar2: _ArrayLike[_EitherSCT],
  342. ) -> NDArray[_EitherSCT]: ...
  343. @overload
  344. def union1d(
  345. ar1: ArrayLike,
  346. ar2: ArrayLike,
  347. ) -> NDArray[Any]: ...
  348. @overload
  349. def setdiff1d(
  350. ar1: _ArrayLike[_EitherSCT],
  351. ar2: _ArrayLike[_EitherSCT],
  352. assume_unique: bool = ...,
  353. ) -> NDArray[_EitherSCT]: ...
  354. @overload
  355. def setdiff1d(
  356. ar1: ArrayLike,
  357. ar2: ArrayLike,
  358. assume_unique: bool = ...,
  359. ) -> NDArray[Any]: ...