hashtable.pyi 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232
  1. from typing import (
  2. Any,
  3. Hashable,
  4. Literal,
  5. )
  6. import numpy as np
  7. def unique_label_indices(
  8. labels: np.ndarray, # const int64_t[:]
  9. ) -> np.ndarray: ...
  10. class Factorizer:
  11. count: int
  12. def __init__(self, size_hint: int): ...
  13. def get_count(self) -> int: ...
  14. class ObjectFactorizer(Factorizer):
  15. table: PyObjectHashTable
  16. uniques: ObjectVector
  17. def factorize(
  18. self,
  19. values: np.ndarray, # ndarray[object]
  20. sort: bool = ...,
  21. na_sentinel=...,
  22. na_value=...,
  23. ) -> np.ndarray: ... # np.ndarray[intp]
  24. class Int64Factorizer(Factorizer):
  25. table: Int64HashTable
  26. uniques: Int64Vector
  27. def factorize(
  28. self,
  29. values: np.ndarray, # const int64_t[:]
  30. sort: bool = ...,
  31. na_sentinel=...,
  32. na_value=...,
  33. ) -> np.ndarray: ... # np.ndarray[intp]
  34. class Int64Vector:
  35. def __init__(self): ...
  36. def __len__(self) -> int: ...
  37. def to_array(self) -> np.ndarray: ... # np.ndarray[np.int64]
  38. class Int32Vector:
  39. def __init__(self): ...
  40. def __len__(self) -> int: ...
  41. def to_array(self) -> np.ndarray: ... # np.ndarray[np.int32]
  42. class Int16Vector:
  43. def __init__(self): ...
  44. def __len__(self) -> int: ...
  45. def to_array(self) -> np.ndarray: ... # np.ndarray[np.int16]
  46. class Int8Vector:
  47. def __init__(self): ...
  48. def __len__(self) -> int: ...
  49. def to_array(self) -> np.ndarray: ... # np.ndarray[np.int8]
  50. class UInt64Vector:
  51. def __init__(self): ...
  52. def __len__(self) -> int: ...
  53. def to_array(self) -> np.ndarray: ... # np.ndarray[np.uint64]
  54. class UInt32Vector:
  55. def __init__(self): ...
  56. def __len__(self) -> int: ...
  57. def to_array(self) -> np.ndarray: ... # np.ndarray[np.uint32]
  58. class UInt16Vector:
  59. def __init__(self): ...
  60. def __len__(self) -> int: ...
  61. def to_array(self) -> np.ndarray: ... # np.ndarray[np.uint16]
  62. class UInt8Vector:
  63. def __init__(self): ...
  64. def __len__(self) -> int: ...
  65. def to_array(self) -> np.ndarray: ... # np.ndarray[np.uint8]
  66. class Float64Vector:
  67. def __init__(self): ...
  68. def __len__(self) -> int: ...
  69. def to_array(self) -> np.ndarray: ... # np.ndarray[np.float64]
  70. class Float32Vector:
  71. def __init__(self): ...
  72. def __len__(self) -> int: ...
  73. def to_array(self) -> np.ndarray: ... # np.ndarray[np.float32]
  74. class Complex128Vector:
  75. def __init__(self): ...
  76. def __len__(self) -> int: ...
  77. def to_array(self) -> np.ndarray: ... # np.ndarray[np.complex128]
  78. class Complex64Vector:
  79. def __init__(self): ...
  80. def __len__(self) -> int: ...
  81. def to_array(self) -> np.ndarray: ... # np.ndarray[np.complex64]
  82. class StringVector:
  83. def __init__(self): ...
  84. def __len__(self) -> int: ...
  85. def to_array(self) -> np.ndarray: ... # np.ndarray[object]
  86. class ObjectVector:
  87. def __init__(self): ...
  88. def __len__(self) -> int: ...
  89. def to_array(self) -> np.ndarray: ... # np.ndarray[object]
  90. class HashTable:
  91. # NB: The base HashTable class does _not_ actually have these methods;
  92. # we are putting the here for the sake of mypy to avoid
  93. # reproducing them in each subclass below.
  94. def __init__(self, size_hint: int = ...): ...
  95. def __len__(self) -> int: ...
  96. def __contains__(self, key: Hashable) -> bool: ...
  97. def sizeof(self, deep: bool = ...) -> int: ...
  98. def get_state(self) -> dict[str, int]: ...
  99. # TODO: `item` type is subclass-specific
  100. def get_item(self, item): ... # TODO: return type?
  101. def set_item(self, item) -> None: ...
  102. # FIXME: we don't actually have this for StringHashTable or ObjectHashTable?
  103. def map(
  104. self,
  105. keys: np.ndarray, # np.ndarray[subclass-specific]
  106. values: np.ndarray, # const int64_t[:]
  107. ) -> None: ...
  108. def map_locations(
  109. self,
  110. values: np.ndarray, # np.ndarray[subclass-specific]
  111. ) -> None: ...
  112. def lookup(
  113. self,
  114. values: np.ndarray, # np.ndarray[subclass-specific]
  115. ) -> np.ndarray: ... # np.ndarray[np.intp]
  116. def get_labels(
  117. self,
  118. values: np.ndarray, # np.ndarray[subclass-specific]
  119. uniques, # SubclassTypeVector
  120. count_prior: int = ...,
  121. na_sentinel: int = ...,
  122. na_value: object = ...,
  123. ) -> np.ndarray: ... # np.ndarray[intp_t]
  124. def unique(
  125. self,
  126. values: np.ndarray, # np.ndarray[subclass-specific]
  127. return_inverse: bool = ...,
  128. ) -> tuple[
  129. np.ndarray, # np.ndarray[subclass-specific]
  130. np.ndarray, # np.ndarray[np.intp],
  131. ] | np.ndarray: ... # np.ndarray[subclass-specific]
  132. def _unique(
  133. self,
  134. values: np.ndarray, # np.ndarray[subclass-specific]
  135. uniques, # FooVector
  136. count_prior: int = ...,
  137. na_sentinel: int = ...,
  138. na_value: object = ...,
  139. ignore_na: bool = ...,
  140. return_inverse: bool = ...,
  141. ) -> tuple[
  142. np.ndarray, # np.ndarray[subclass-specific]
  143. np.ndarray, # np.ndarray[np.intp],
  144. ] | np.ndarray: ... # np.ndarray[subclass-specific]
  145. def factorize(
  146. self,
  147. values: np.ndarray, # np.ndarray[subclass-specific]
  148. na_sentinel: int = ...,
  149. na_value: object = ...,
  150. mask=...,
  151. ) -> tuple[
  152. np.ndarray, # np.ndarray[subclass-specific]
  153. np.ndarray, # np.ndarray[np.intp],
  154. ]: ...
  155. class Complex128HashTable(HashTable): ...
  156. class Complex64HashTable(HashTable): ...
  157. class Float64HashTable(HashTable): ...
  158. class Float32HashTable(HashTable): ...
  159. class Int64HashTable(HashTable):
  160. # Only Int64HashTable has get_labels_groupby
  161. def get_labels_groupby(
  162. self,
  163. values: np.ndarray, # const int64_t[:]
  164. ) -> tuple[
  165. np.ndarray, # np.ndarray[np.intp]
  166. np.ndarray, # np.ndarray[np.int64]
  167. ]: ...
  168. class Int32HashTable(HashTable): ...
  169. class Int16HashTable(HashTable): ...
  170. class Int8HashTable(HashTable): ...
  171. class UInt64HashTable(HashTable): ...
  172. class UInt32HashTable(HashTable): ...
  173. class UInt16HashTable(HashTable): ...
  174. class UInt8HashTable(HashTable): ...
  175. class StringHashTable(HashTable): ...
  176. class PyObjectHashTable(HashTable): ...
  177. def duplicated_int64(
  178. values: np.ndarray, # const int64_t[:] values
  179. keep: Literal["last", "first", False] = ...,
  180. ) -> np.ndarray: ... # np.ndarray[bool]
  181. # TODO: Is it actually bool or is it uint8?
  182. def mode_int64(
  183. values: np.ndarray, # const int64_t[:] values
  184. dropna: bool,
  185. ) -> np.ndarray: ... # np.ndarray[np.int64]
  186. def value_count_int64(
  187. values: np.ndarray, # const int64_t[:]
  188. dropna: bool,
  189. ) -> tuple[np.ndarray, np.ndarray,]: ... # np.ndarray[np.int64] # np.ndarray[np.int64]
  190. def duplicated(
  191. values: np.ndarray,
  192. keep: Literal["last", "first", False] = ...,
  193. ) -> np.ndarray: ... # np.ndarray[bool]
  194. def mode(values: np.ndarray, dropna: bool) -> np.ndarray: ...
  195. def value_count(
  196. values: np.ndarray,
  197. dropna: bool,
  198. ) -> tuple[np.ndarray, np.ndarray,]: ... # np.ndarray[np.int64]
  199. # arr and values should have same dtype
  200. def ismember(
  201. arr: np.ndarray,
  202. values: np.ndarray,
  203. ) -> np.ndarray: ... # np.ndarray[bool]
  204. def object_hash(obj) -> int: ...
  205. def objects_are_equal(a, b) -> bool: ...