arraysetops.pyi 8.2 KB

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