lib.pyi 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238
  1. # TODO(npdtypes): Many types specified here can be made more specific/accurate;
  2. # the more specific versions are specified in comments
  3. from collections.abc import (
  4. Callable,
  5. Generator,
  6. Hashable,
  7. )
  8. from decimal import Decimal
  9. from typing import (
  10. Any,
  11. Final,
  12. Literal,
  13. TypeAlias,
  14. TypeGuard,
  15. overload,
  16. )
  17. import numpy as np
  18. from pandas._typing import (
  19. ArrayLike,
  20. DtypeObj,
  21. npt,
  22. )
  23. # placeholder until we can specify np.ndarray[object, ndim=2]
  24. ndarray_obj_2d = np.ndarray
  25. from enum import Enum
  26. class _NoDefault(Enum):
  27. no_default = ...
  28. no_default: Final = _NoDefault.no_default
  29. NoDefault: TypeAlias = Literal[_NoDefault.no_default]
  30. i8max: int
  31. u8max: int
  32. def is_np_dtype(dtype: object, kinds: str | None = ...) -> TypeGuard[np.dtype]: ...
  33. def item_from_zerodim(val: object) -> object: ...
  34. def infer_dtype(value: object, skipna: bool = ...) -> str: ...
  35. def is_iterator(obj: object) -> bool: ...
  36. def is_scalar(val: object) -> bool: ...
  37. def is_list_like(obj: object, allow_sets: bool = ...) -> bool: ...
  38. def is_pyarrow_array(obj: object) -> bool: ...
  39. def is_decimal(obj: object) -> TypeGuard[Decimal]: ...
  40. def is_complex(obj: object) -> TypeGuard[complex]: ...
  41. def is_bool(obj: object) -> TypeGuard[bool | np.bool_]: ...
  42. def is_integer(obj: object) -> TypeGuard[int | np.integer]: ...
  43. def is_int_or_none(obj) -> bool: ...
  44. def is_float(obj: object) -> TypeGuard[float]: ...
  45. def is_interval_array(values: np.ndarray) -> bool: ...
  46. def is_datetime64_array(values: np.ndarray, skipna: bool = True) -> bool: ...
  47. def is_timedelta_or_timedelta64_array(
  48. values: np.ndarray, skipna: bool = True
  49. ) -> bool: ...
  50. def is_datetime_with_singletz_array(values: np.ndarray) -> bool: ...
  51. def is_time_array(values: np.ndarray, skipna: bool = ...): ...
  52. def is_date_array(values: np.ndarray, skipna: bool = ...): ...
  53. def is_datetime_array(values: np.ndarray, skipna: bool = ...): ...
  54. def is_string_array(values: np.ndarray, skipna: bool = ...): ...
  55. def is_float_array(values: np.ndarray, skipna: bool = ...): ...
  56. def is_integer_array(values: np.ndarray, skipna: bool = ...): ...
  57. def is_bool_array(values: np.ndarray, skipna: bool = ...): ...
  58. def fast_multiget(
  59. mapping: dict,
  60. keys: np.ndarray, # object[:]
  61. default=...,
  62. ) -> ArrayLike: ...
  63. def fast_unique_multiple_list_gen(gen: Generator, sort: bool = ...) -> list: ...
  64. @overload
  65. def map_infer(
  66. arr: np.ndarray,
  67. f: Callable[[Any], Any],
  68. *,
  69. convert: Literal[False],
  70. ignore_na: bool = ...,
  71. ) -> np.ndarray: ...
  72. @overload
  73. def map_infer(
  74. arr: np.ndarray,
  75. f: Callable[[Any], Any],
  76. *,
  77. convert: bool = ...,
  78. ignore_na: bool = ...,
  79. ) -> ArrayLike: ...
  80. @overload
  81. def maybe_convert_objects(
  82. objects: npt.NDArray[np.object_],
  83. *,
  84. try_float: bool = ...,
  85. safe: bool = ...,
  86. convert_numeric: bool = ...,
  87. convert_non_numeric: Literal[False] = ...,
  88. convert_to_nullable_dtype: Literal[False] = ...,
  89. dtype_if_all_nat: DtypeObj | None = ...,
  90. ) -> npt.NDArray[np.object_ | np.number]: ...
  91. @overload
  92. def maybe_convert_objects(
  93. objects: npt.NDArray[np.object_],
  94. *,
  95. try_float: bool = ...,
  96. safe: bool = ...,
  97. convert_numeric: bool = ...,
  98. convert_non_numeric: bool = ...,
  99. convert_to_nullable_dtype: Literal[True] = ...,
  100. dtype_if_all_nat: DtypeObj | None = ...,
  101. ) -> ArrayLike: ...
  102. @overload
  103. def maybe_convert_objects(
  104. objects: npt.NDArray[np.object_],
  105. *,
  106. try_float: bool = ...,
  107. safe: bool = ...,
  108. convert_numeric: bool = ...,
  109. convert_non_numeric: bool = ...,
  110. convert_to_nullable_dtype: bool = ...,
  111. dtype_if_all_nat: DtypeObj | None = ...,
  112. ) -> ArrayLike: ...
  113. @overload
  114. def maybe_convert_numeric(
  115. values: npt.NDArray[np.object_],
  116. na_values: set,
  117. convert_empty: bool = ...,
  118. coerce_numeric: bool = ...,
  119. convert_to_masked_nullable: Literal[False] = ...,
  120. ) -> tuple[np.ndarray, None]: ...
  121. @overload
  122. def maybe_convert_numeric(
  123. values: npt.NDArray[np.object_],
  124. na_values: set,
  125. convert_empty: bool = ...,
  126. coerce_numeric: bool = ...,
  127. *,
  128. convert_to_masked_nullable: Literal[True],
  129. ) -> tuple[np.ndarray, np.ndarray]: ...
  130. # TODO: restrict `arr`?
  131. def ensure_string_array(
  132. arr,
  133. na_value: object = ...,
  134. convert_na_value: bool = ...,
  135. copy: bool = ...,
  136. skipna: bool = ...,
  137. ) -> npt.NDArray[np.object_]: ...
  138. def convert_nans_to_NA(
  139. arr: npt.NDArray[np.object_],
  140. ) -> npt.NDArray[np.object_]: ...
  141. def fast_zip(ndarrays: list) -> npt.NDArray[np.object_]: ...
  142. # TODO: can we be more specific about rows?
  143. def to_object_array_tuples(rows: object) -> ndarray_obj_2d: ...
  144. def tuples_to_object_array(
  145. tuples: npt.NDArray[np.object_],
  146. ) -> ndarray_obj_2d: ...
  147. # TODO: can we be more specific about rows?
  148. def to_object_array(rows: object, min_width: int = ...) -> ndarray_obj_2d: ...
  149. def dicts_to_array(dicts: list, columns: list) -> ndarray_obj_2d: ...
  150. def maybe_booleans_to_slice(
  151. mask: npt.NDArray[np.uint8],
  152. ) -> slice | npt.NDArray[np.uint8]: ...
  153. def maybe_indices_to_slice(
  154. indices: npt.NDArray[np.intp],
  155. max_len: int,
  156. ) -> slice | npt.NDArray[np.intp]: ...
  157. def is_all_arraylike(obj: list) -> bool: ...
  158. # -----------------------------------------------------------------
  159. # Functions which in reality take memoryviews
  160. def memory_usage_of_objects(arr: np.ndarray) -> int: ... # object[:] # np.int64
  161. @overload
  162. def map_infer_mask(
  163. arr: np.ndarray,
  164. f: Callable[[Any], Any],
  165. mask: np.ndarray, # const uint8_t[:]
  166. *,
  167. convert: Literal[False],
  168. na_value: Any = ...,
  169. dtype: np.dtype = ...,
  170. ) -> np.ndarray: ...
  171. @overload
  172. def map_infer_mask(
  173. arr: np.ndarray,
  174. f: Callable[[Any], Any],
  175. mask: np.ndarray, # const uint8_t[:]
  176. *,
  177. convert: bool = ...,
  178. na_value: Any = ...,
  179. dtype: np.dtype = ...,
  180. ) -> ArrayLike: ...
  181. def indices_fast(
  182. index: npt.NDArray[np.intp],
  183. labels: np.ndarray, # const int64_t[:]
  184. keys: list,
  185. sorted_labels: list[npt.NDArray[np.int64]],
  186. ) -> dict[Hashable, npt.NDArray[np.intp]]: ...
  187. def generate_slices(
  188. labels: np.ndarray,
  189. ngroups: int, # const intp_t[:]
  190. ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]: ...
  191. def count_level_2d(
  192. mask: np.ndarray, # ndarray[uint8_t, ndim=2, cast=True],
  193. labels: np.ndarray, # const intp_t[:]
  194. max_bin: int,
  195. ) -> np.ndarray: ... # np.ndarray[np.int64, ndim=2]
  196. def get_level_sorter(
  197. codes: np.ndarray, # const int64_t[:]
  198. starts: np.ndarray, # const intp_t[:]
  199. ) -> np.ndarray: ... # np.ndarray[np.intp, ndim=1]
  200. def generate_bins_dt64(
  201. values: npt.NDArray[np.int64],
  202. binner: np.ndarray, # const int64_t[:]
  203. closed: object = ...,
  204. hasnans: bool = ...,
  205. ) -> np.ndarray: ... # np.ndarray[np.int64, ndim=1]
  206. def array_equivalent_object(
  207. left: npt.NDArray[np.object_],
  208. right: npt.NDArray[np.object_],
  209. ) -> bool: ...
  210. def has_infs(arr: np.ndarray) -> bool: ... # const floating[:]
  211. def has_only_ints_or_nan(arr: np.ndarray) -> bool: ... # const floating[:]
  212. def get_reverse_indexer(
  213. indexer: np.ndarray, # const intp_t[:]
  214. length: int,
  215. ) -> npt.NDArray[np.intp]: ...
  216. def is_bool_list(obj: list) -> bool: ...
  217. def dtypes_all_equal(types: list[DtypeObj]) -> bool: ...
  218. def is_range_indexer(
  219. left: np.ndarray,
  220. n: int, # np.ndarray[np.int64, ndim=1]
  221. ) -> bool: ...
  222. def is_sequence_range(
  223. sequence: np.ndarray,
  224. step: int, # np.ndarray[np.int64, ndim=1]
  225. ) -> bool: ...