recfunctions.pyi 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435
  1. from collections.abc import Callable, Iterable, Mapping, Sequence
  2. from typing import Any, Literal, TypeAlias, overload
  3. from _typeshed import Incomplete
  4. from typing_extensions import TypeVar
  5. import numpy as np
  6. import numpy.typing as npt
  7. from numpy._typing import _DTypeLike, _DTypeLikeVoid
  8. from numpy.ma.mrecords import MaskedRecords
  9. __all__ = [
  10. "append_fields",
  11. "apply_along_fields",
  12. "assign_fields_by_name",
  13. "drop_fields",
  14. "find_duplicates",
  15. "flatten_descr",
  16. "get_fieldstructure",
  17. "get_names",
  18. "get_names_flat",
  19. "join_by",
  20. "merge_arrays",
  21. "rec_append_fields",
  22. "rec_drop_fields",
  23. "rec_join",
  24. "recursive_fill_fields",
  25. "rename_fields",
  26. "repack_fields",
  27. "require_fields",
  28. "stack_arrays",
  29. "structured_to_unstructured",
  30. "unstructured_to_structured",
  31. ]
  32. _T = TypeVar("_T")
  33. _ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...])
  34. _ScalarT = TypeVar("_ScalarT", bound=np.generic)
  35. _DTypeT = TypeVar("_DTypeT", bound=np.dtype[Any])
  36. _ArrayT = TypeVar("_ArrayT", bound=npt.NDArray[Any])
  37. _VoidArrayT = TypeVar("_VoidArrayT", bound=npt.NDArray[np.void])
  38. _NonVoidDTypeT = TypeVar("_NonVoidDTypeT", bound=_NonVoidDType)
  39. _OneOrMany: TypeAlias = _T | Iterable[_T]
  40. _BuiltinSequence: TypeAlias = tuple[_T, ...] | list[_T]
  41. _NestedNames: TypeAlias = tuple[str | _NestedNames, ...]
  42. _NonVoid: TypeAlias = np.bool | np.number | np.character | np.datetime64 | np.timedelta64 | np.object_
  43. _NonVoidDType: TypeAlias = np.dtype[_NonVoid] | np.dtypes.StringDType
  44. _JoinType: TypeAlias = Literal["inner", "outer", "leftouter"]
  45. ###
  46. def recursive_fill_fields(input: npt.NDArray[np.void], output: _VoidArrayT) -> _VoidArrayT: ...
  47. #
  48. def get_names(adtype: np.dtype[np.void]) -> _NestedNames: ...
  49. def get_names_flat(adtype: np.dtype[np.void]) -> tuple[str, ...]: ...
  50. #
  51. @overload
  52. def flatten_descr(ndtype: _NonVoidDTypeT) -> tuple[tuple[Literal[""], _NonVoidDTypeT]]: ...
  53. @overload
  54. def flatten_descr(ndtype: np.dtype[np.void]) -> tuple[tuple[str, np.dtype[Any]]]: ...
  55. #
  56. def get_fieldstructure(
  57. adtype: np.dtype[np.void],
  58. lastname: str | None = None,
  59. parents: dict[str, list[str]] | None = None,
  60. ) -> dict[str, list[str]]: ...
  61. #
  62. @overload
  63. def merge_arrays(
  64. seqarrays: Sequence[np.ndarray[_ShapeT, np.dtype[Any]]] | np.ndarray[_ShapeT, np.dtype[Any]],
  65. fill_value: float = -1,
  66. flatten: bool = False,
  67. usemask: bool = False,
  68. asrecarray: bool = False,
  69. ) -> np.recarray[_ShapeT, np.dtype[np.void]]: ...
  70. @overload
  71. def merge_arrays(
  72. seqarrays: Sequence[npt.ArrayLike] | np.void,
  73. fill_value: float = -1,
  74. flatten: bool = False,
  75. usemask: bool = False,
  76. asrecarray: bool = False,
  77. ) -> np.recarray[Any, np.dtype[np.void]]: ...
  78. #
  79. @overload
  80. def drop_fields(
  81. base: np.ndarray[_ShapeT, np.dtype[np.void]],
  82. drop_names: str | Iterable[str],
  83. usemask: bool = True,
  84. asrecarray: Literal[False] = False,
  85. ) -> np.ndarray[_ShapeT, np.dtype[np.void]]: ...
  86. @overload
  87. def drop_fields(
  88. base: np.ndarray[_ShapeT, np.dtype[np.void]],
  89. drop_names: str | Iterable[str],
  90. usemask: bool,
  91. asrecarray: Literal[True],
  92. ) -> np.recarray[_ShapeT, np.dtype[np.void]]: ...
  93. @overload
  94. def drop_fields(
  95. base: np.ndarray[_ShapeT, np.dtype[np.void]],
  96. drop_names: str | Iterable[str],
  97. usemask: bool = True,
  98. *,
  99. asrecarray: Literal[True],
  100. ) -> np.recarray[_ShapeT, np.dtype[np.void]]: ...
  101. #
  102. @overload
  103. def rename_fields(
  104. base: MaskedRecords[_ShapeT, np.dtype[np.void]],
  105. namemapper: Mapping[str, str],
  106. ) -> MaskedRecords[_ShapeT, np.dtype[np.void]]: ...
  107. @overload
  108. def rename_fields(
  109. base: np.ma.MaskedArray[_ShapeT, np.dtype[np.void]],
  110. namemapper: Mapping[str, str],
  111. ) -> np.ma.MaskedArray[_ShapeT, np.dtype[np.void]]: ...
  112. @overload
  113. def rename_fields(
  114. base: np.recarray[_ShapeT, np.dtype[np.void]],
  115. namemapper: Mapping[str, str],
  116. ) -> np.recarray[_ShapeT, np.dtype[np.void]]: ...
  117. @overload
  118. def rename_fields(
  119. base: np.ndarray[_ShapeT, np.dtype[np.void]],
  120. namemapper: Mapping[str, str],
  121. ) -> np.ndarray[_ShapeT, np.dtype[np.void]]: ...
  122. #
  123. @overload
  124. def append_fields(
  125. base: np.ndarray[_ShapeT, np.dtype[np.void]],
  126. names: _OneOrMany[str],
  127. data: _OneOrMany[npt.NDArray[Any]],
  128. dtypes: _BuiltinSequence[np.dtype[Any]] | None,
  129. fill_value: int,
  130. usemask: Literal[False],
  131. asrecarray: Literal[False] = False,
  132. ) -> np.ndarray[_ShapeT, np.dtype[np.void]]: ...
  133. @overload
  134. def append_fields(
  135. base: np.ndarray[_ShapeT, np.dtype[np.void]],
  136. names: _OneOrMany[str],
  137. data: _OneOrMany[npt.NDArray[Any]],
  138. dtypes: _BuiltinSequence[np.dtype[Any]] | None = None,
  139. fill_value: int = -1,
  140. *,
  141. usemask: Literal[False],
  142. asrecarray: Literal[False] = False,
  143. ) -> np.ndarray[_ShapeT, np.dtype[np.void]]: ...
  144. @overload
  145. def append_fields(
  146. base: np.ndarray[_ShapeT, np.dtype[np.void]],
  147. names: _OneOrMany[str],
  148. data: _OneOrMany[npt.NDArray[Any]],
  149. dtypes: _BuiltinSequence[np.dtype[Any]] | None,
  150. fill_value: int,
  151. usemask: Literal[False],
  152. asrecarray: Literal[True],
  153. ) -> np.recarray[_ShapeT, np.dtype[np.void]]: ...
  154. @overload
  155. def append_fields(
  156. base: np.ndarray[_ShapeT, np.dtype[np.void]],
  157. names: _OneOrMany[str],
  158. data: _OneOrMany[npt.NDArray[Any]],
  159. dtypes: _BuiltinSequence[np.dtype[Any]] | None = None,
  160. fill_value: int = -1,
  161. *,
  162. usemask: Literal[False],
  163. asrecarray: Literal[True],
  164. ) -> np.recarray[_ShapeT, np.dtype[np.void]]: ...
  165. @overload
  166. def append_fields(
  167. base: np.ndarray[_ShapeT, np.dtype[np.void]],
  168. names: _OneOrMany[str],
  169. data: _OneOrMany[npt.NDArray[Any]],
  170. dtypes: _BuiltinSequence[np.dtype[Any]] | None = None,
  171. fill_value: int = -1,
  172. usemask: Literal[True] = True,
  173. asrecarray: Literal[False] = False,
  174. ) -> np.ma.MaskedArray[_ShapeT, np.dtype[np.void]]: ...
  175. @overload
  176. def append_fields(
  177. base: np.ndarray[_ShapeT, np.dtype[np.void]],
  178. names: _OneOrMany[str],
  179. data: _OneOrMany[npt.NDArray[Any]],
  180. dtypes: _BuiltinSequence[np.dtype[Any]] | None,
  181. fill_value: int,
  182. usemask: Literal[True],
  183. asrecarray: Literal[True],
  184. ) -> MaskedRecords[_ShapeT, np.dtype[np.void]]: ...
  185. @overload
  186. def append_fields(
  187. base: np.ndarray[_ShapeT, np.dtype[np.void]],
  188. names: _OneOrMany[str],
  189. data: _OneOrMany[npt.NDArray[Any]],
  190. dtypes: _BuiltinSequence[np.dtype[Any]] | None = None,
  191. fill_value: int = -1,
  192. usemask: Literal[True] = True,
  193. *,
  194. asrecarray: Literal[True],
  195. ) -> MaskedRecords[_ShapeT, np.dtype[np.void]]: ...
  196. #
  197. def rec_drop_fields(
  198. base: np.ndarray[_ShapeT, np.dtype[np.void]],
  199. drop_names: str | Iterable[str],
  200. ) -> np.recarray[_ShapeT, np.dtype[np.void]]: ...
  201. #
  202. def rec_append_fields(
  203. base: np.ndarray[_ShapeT, np.dtype[np.void]],
  204. names: _OneOrMany[str],
  205. data: _OneOrMany[npt.NDArray[Any]],
  206. dtypes: _BuiltinSequence[np.dtype[Any]] | None = None,
  207. ) -> np.ma.MaskedArray[_ShapeT, np.dtype[np.void]]: ...
  208. # TODO(jorenham): Stop passing `void` directly once structured dtypes are implemented,
  209. # e.g. using a `TypeVar` with constraints.
  210. # https://github.com/numpy/numtype/issues/92
  211. @overload
  212. def repack_fields(a: _DTypeT, align: bool = False, recurse: bool = False) -> _DTypeT: ...
  213. @overload
  214. def repack_fields(a: _ScalarT, align: bool = False, recurse: bool = False) -> _ScalarT: ...
  215. @overload
  216. def repack_fields(a: _ArrayT, align: bool = False, recurse: bool = False) -> _ArrayT: ...
  217. # TODO(jorenham): Attempt shape-typing (return type has ndim == arr.ndim + 1)
  218. @overload
  219. def structured_to_unstructured(
  220. arr: npt.NDArray[np.void],
  221. dtype: _DTypeLike[_ScalarT],
  222. copy: bool = False,
  223. casting: np._CastingKind = "unsafe",
  224. ) -> npt.NDArray[_ScalarT]: ...
  225. @overload
  226. def structured_to_unstructured(
  227. arr: npt.NDArray[np.void],
  228. dtype: npt.DTypeLike | None = None,
  229. copy: bool = False,
  230. casting: np._CastingKind = "unsafe",
  231. ) -> npt.NDArray[Any]: ...
  232. #
  233. @overload
  234. def unstructured_to_structured(
  235. arr: npt.NDArray[Any],
  236. dtype: npt.DTypeLike,
  237. names: None = None,
  238. align: bool = False,
  239. copy: bool = False,
  240. casting: str = "unsafe",
  241. ) -> npt.NDArray[np.void]: ...
  242. @overload
  243. def unstructured_to_structured(
  244. arr: npt.NDArray[Any],
  245. dtype: None,
  246. names: _OneOrMany[str],
  247. align: bool = False,
  248. copy: bool = False,
  249. casting: str = "unsafe",
  250. ) -> npt.NDArray[np.void]: ...
  251. #
  252. def apply_along_fields(
  253. func: Callable[[np.ndarray[_ShapeT, Any]], npt.NDArray[Any]],
  254. arr: np.ndarray[_ShapeT, np.dtype[np.void]],
  255. ) -> np.ndarray[_ShapeT, np.dtype[np.void]]: ...
  256. #
  257. def assign_fields_by_name(dst: npt.NDArray[np.void], src: npt.NDArray[np.void], zero_unassigned: bool = True) -> None: ...
  258. #
  259. def require_fields(
  260. array: np.ndarray[_ShapeT, np.dtype[np.void]],
  261. required_dtype: _DTypeLikeVoid,
  262. ) -> np.ndarray[_ShapeT, np.dtype[np.void]]: ...
  263. # TODO(jorenham): Attempt shape-typing
  264. @overload
  265. def stack_arrays(
  266. arrays: _ArrayT,
  267. defaults: Mapping[str, object] | None = None,
  268. usemask: bool = True,
  269. asrecarray: bool = False,
  270. autoconvert: bool = False,
  271. ) -> _ArrayT: ...
  272. @overload
  273. def stack_arrays(
  274. arrays: Sequence[npt.NDArray[Any]],
  275. defaults: Mapping[str, Incomplete] | None,
  276. usemask: Literal[False],
  277. asrecarray: Literal[False] = False,
  278. autoconvert: bool = False,
  279. ) -> npt.NDArray[np.void]: ...
  280. @overload
  281. def stack_arrays(
  282. arrays: Sequence[npt.NDArray[Any]],
  283. defaults: Mapping[str, Incomplete] | None = None,
  284. *,
  285. usemask: Literal[False],
  286. asrecarray: Literal[False] = False,
  287. autoconvert: bool = False,
  288. ) -> npt.NDArray[np.void]: ...
  289. @overload
  290. def stack_arrays(
  291. arrays: Sequence[npt.NDArray[Any]],
  292. defaults: Mapping[str, Incomplete] | None = None,
  293. *,
  294. usemask: Literal[False],
  295. asrecarray: Literal[True],
  296. autoconvert: bool = False,
  297. ) -> np.recarray[tuple[int, ...], np.dtype[np.void]]: ...
  298. @overload
  299. def stack_arrays(
  300. arrays: Sequence[npt.NDArray[Any]],
  301. defaults: Mapping[str, Incomplete] | None = None,
  302. usemask: Literal[True] = True,
  303. asrecarray: Literal[False] = False,
  304. autoconvert: bool = False,
  305. ) -> np.ma.MaskedArray[tuple[int, ...], np.dtype[np.void]]: ...
  306. @overload
  307. def stack_arrays(
  308. arrays: Sequence[npt.NDArray[Any]],
  309. defaults: Mapping[str, Incomplete] | None,
  310. usemask: Literal[True],
  311. asrecarray: Literal[True],
  312. autoconvert: bool = False,
  313. ) -> MaskedRecords[tuple[int, ...], np.dtype[np.void]]: ...
  314. @overload
  315. def stack_arrays(
  316. arrays: Sequence[npt.NDArray[Any]],
  317. defaults: Mapping[str, Incomplete] | None = None,
  318. usemask: Literal[True] = True,
  319. *,
  320. asrecarray: Literal[True],
  321. autoconvert: bool = False,
  322. ) -> MaskedRecords[tuple[int, ...], np.dtype[np.void]]: ...
  323. #
  324. @overload
  325. def find_duplicates(
  326. a: np.ma.MaskedArray[_ShapeT, np.dtype[np.void]],
  327. key: str | None = None,
  328. ignoremask: bool = True,
  329. return_index: Literal[False] = False,
  330. ) -> np.ma.MaskedArray[_ShapeT, np.dtype[np.void]]: ...
  331. @overload
  332. def find_duplicates(
  333. a: np.ma.MaskedArray[_ShapeT, np.dtype[np.void]],
  334. key: str | None,
  335. ignoremask: bool,
  336. return_index: Literal[True],
  337. ) -> tuple[np.ma.MaskedArray[_ShapeT, np.dtype[np.void]], np.ndarray[_ShapeT, np.dtype[np.int_]]]: ...
  338. @overload
  339. def find_duplicates(
  340. a: np.ma.MaskedArray[_ShapeT, np.dtype[np.void]],
  341. key: str | None = None,
  342. ignoremask: bool = True,
  343. *,
  344. return_index: Literal[True],
  345. ) -> tuple[np.ma.MaskedArray[_ShapeT, np.dtype[np.void]], np.ndarray[_ShapeT, np.dtype[np.int_]]]: ...
  346. #
  347. @overload
  348. def join_by(
  349. key: str | Sequence[str],
  350. r1: npt.NDArray[np.void],
  351. r2: npt.NDArray[np.void],
  352. jointype: _JoinType = "inner",
  353. r1postfix: str = "1",
  354. r2postfix: str = "2",
  355. defaults: Mapping[str, object] | None = None,
  356. *,
  357. usemask: Literal[False],
  358. asrecarray: Literal[False] = False,
  359. ) -> np.ndarray[tuple[int], np.dtype[np.void]]: ...
  360. @overload
  361. def join_by(
  362. key: str | Sequence[str],
  363. r1: npt.NDArray[np.void],
  364. r2: npt.NDArray[np.void],
  365. jointype: _JoinType = "inner",
  366. r1postfix: str = "1",
  367. r2postfix: str = "2",
  368. defaults: Mapping[str, object] | None = None,
  369. *,
  370. usemask: Literal[False],
  371. asrecarray: Literal[True],
  372. ) -> np.recarray[tuple[int], np.dtype[np.void]]: ...
  373. @overload
  374. def join_by(
  375. key: str | Sequence[str],
  376. r1: npt.NDArray[np.void],
  377. r2: npt.NDArray[np.void],
  378. jointype: _JoinType = "inner",
  379. r1postfix: str = "1",
  380. r2postfix: str = "2",
  381. defaults: Mapping[str, object] | None = None,
  382. usemask: Literal[True] = True,
  383. asrecarray: Literal[False] = False,
  384. ) -> np.ma.MaskedArray[tuple[int], np.dtype[np.void]]: ...
  385. @overload
  386. def join_by(
  387. key: str | Sequence[str],
  388. r1: npt.NDArray[np.void],
  389. r2: npt.NDArray[np.void],
  390. jointype: _JoinType = "inner",
  391. r1postfix: str = "1",
  392. r2postfix: str = "2",
  393. defaults: Mapping[str, object] | None = None,
  394. usemask: Literal[True] = True,
  395. *,
  396. asrecarray: Literal[True],
  397. ) -> MaskedRecords[tuple[int], np.dtype[np.void]]: ...
  398. #
  399. def rec_join(
  400. key: str | Sequence[str],
  401. r1: npt.NDArray[np.void],
  402. r2: npt.NDArray[np.void],
  403. jointype: _JoinType = "inner",
  404. r1postfix: str = "1",
  405. r2postfix: str = "2",
  406. defaults: Mapping[str, object] | None = None,
  407. ) -> np.recarray[tuple[int], np.dtype[np.void]]: ...