You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

298 lines
9.8 KiB

4 years ago
  1. # cython: infer_types=True
  2. # coding: utf8
  3. from __future__ import unicode_literals, absolute_import
  4. cimport cython
  5. from libc.string cimport memcpy
  6. from libcpp.set cimport set
  7. from libc.stdint cimport uint32_t
  8. from murmurhash.mrmr cimport hash64, hash32
  9. import ujson
  10. from .symbols import IDS as SYMBOLS_BY_STR
  11. from .symbols import NAMES as SYMBOLS_BY_INT
  12. from .typedefs cimport hash_t
  13. from .compat import json_dumps
  14. from .errors import Errors
  15. from . import util
  16. cpdef hash_t hash_string(unicode string) except 0:
  17. chars = string.encode('utf8')
  18. return hash_utf8(chars, len(chars))
  19. cdef hash_t hash_utf8(char* utf8_string, int length) nogil:
  20. return hash64(utf8_string, length, 1)
  21. cdef uint32_t hash32_utf8(char* utf8_string, int length) nogil:
  22. return hash32(utf8_string, length, 1)
  23. cdef unicode decode_Utf8Str(const Utf8Str* string):
  24. cdef int i, length
  25. if string.s[0] < sizeof(string.s) and string.s[0] != 0:
  26. return string.s[1:string.s[0]+1].decode('utf8')
  27. elif string.p[0] < 255:
  28. return string.p[1:string.p[0]+1].decode('utf8')
  29. else:
  30. i = 0
  31. length = 0
  32. while string.p[i] == 255:
  33. i += 1
  34. length += 255
  35. length += string.p[i]
  36. i += 1
  37. return string.p[i:length + i].decode('utf8')
  38. cdef Utf8Str* _allocate(Pool mem, const unsigned char* chars, uint32_t length) except *:
  39. cdef int n_length_bytes
  40. cdef int i
  41. cdef Utf8Str* string = <Utf8Str*>mem.alloc(1, sizeof(Utf8Str))
  42. cdef uint32_t ulength = length
  43. if length < sizeof(string.s):
  44. string.s[0] = <unsigned char>length
  45. memcpy(&string.s[1], chars, length)
  46. return string
  47. elif length < 255:
  48. string.p = <unsigned char*>mem.alloc(length + 1, sizeof(unsigned char))
  49. string.p[0] = length
  50. memcpy(&string.p[1], chars, length)
  51. return string
  52. else:
  53. i = 0
  54. n_length_bytes = (length // 255) + 1
  55. string.p = <unsigned char*>mem.alloc(length + n_length_bytes, sizeof(unsigned char))
  56. for i in range(n_length_bytes-1):
  57. string.p[i] = 255
  58. string.p[n_length_bytes-1] = length % 255
  59. memcpy(&string.p[n_length_bytes], chars, length)
  60. return string
  61. cdef class StringStore:
  62. """Look up strings by 64-bit hashes."""
  63. def __init__(self, strings=None, freeze=False):
  64. """Create the StringStore.
  65. strings (iterable): A sequence of unicode strings to add to the store.
  66. RETURNS (StringStore): The newly constructed object.
  67. """
  68. self.mem = Pool()
  69. self._map = PreshMap()
  70. if strings is not None:
  71. for string in strings:
  72. self.add(string)
  73. def __getitem__(self, object string_or_id):
  74. """Retrieve a string from a given hash, or vice versa.
  75. string_or_id (bytes, unicode or uint64): The value to encode.
  76. Returns (unicode or uint64): The value to be retrieved.
  77. """
  78. if isinstance(string_or_id, basestring) and len(string_or_id) == 0:
  79. return 0
  80. elif string_or_id == 0:
  81. return u''
  82. elif string_or_id in SYMBOLS_BY_STR:
  83. return SYMBOLS_BY_STR[string_or_id]
  84. cdef hash_t key
  85. if isinstance(string_or_id, unicode):
  86. key = hash_string(string_or_id)
  87. return key
  88. elif isinstance(string_or_id, bytes):
  89. key = hash_utf8(string_or_id, len(string_or_id))
  90. return key
  91. elif string_or_id < len(SYMBOLS_BY_INT):
  92. return SYMBOLS_BY_INT[string_or_id]
  93. else:
  94. key = string_or_id
  95. self.hits.insert(key)
  96. utf8str = <Utf8Str*>self._map.get(key)
  97. if utf8str is NULL:
  98. raise KeyError(Errors.E018.format(hash_value=string_or_id))
  99. else:
  100. return decode_Utf8Str(utf8str)
  101. def add(self, string):
  102. """Add a string to the StringStore.
  103. string (unicode): The string to add.
  104. RETURNS (uint64): The string's hash value.
  105. """
  106. if isinstance(string, unicode):
  107. if string in SYMBOLS_BY_STR:
  108. return SYMBOLS_BY_STR[string]
  109. key = hash_string(string)
  110. self.intern_unicode(string)
  111. elif isinstance(string, bytes):
  112. if string in SYMBOLS_BY_STR:
  113. return SYMBOLS_BY_STR[string]
  114. key = hash_utf8(string, len(string))
  115. self._intern_utf8(string, len(string))
  116. else:
  117. raise TypeError(Errors.E017.format(value_type=type(string)))
  118. return key
  119. def __len__(self):
  120. """The number of strings in the store.
  121. RETURNS (int): The number of strings in the store.
  122. """
  123. return self.keys.size()
  124. def __contains__(self, string not None):
  125. """Check whether a string is in the store.
  126. string (unicode): The string to check.
  127. RETURNS (bool): Whether the store contains the string.
  128. """
  129. cdef hash_t key
  130. if isinstance(string, int) or isinstance(string, long):
  131. if string == 0:
  132. return True
  133. key = string
  134. elif len(string) == 0:
  135. return True
  136. elif string in SYMBOLS_BY_STR:
  137. return True
  138. elif isinstance(string, unicode):
  139. key = hash_string(string)
  140. else:
  141. string = string.encode('utf8')
  142. key = hash_utf8(string, len(string))
  143. if key < len(SYMBOLS_BY_INT):
  144. return True
  145. else:
  146. self.hits.insert(key)
  147. return self._map.get(key) is not NULL
  148. def __iter__(self):
  149. """Iterate over the strings in the store, in order.
  150. YIELDS (unicode): A string in the store.
  151. """
  152. cdef int i
  153. cdef hash_t key
  154. for i in range(self.keys.size()):
  155. key = self.keys[i]
  156. self.hits.insert(key)
  157. utf8str = <Utf8Str*>self._map.get(key)
  158. yield decode_Utf8Str(utf8str)
  159. # TODO: Iterate OOV here?
  160. def __reduce__(self):
  161. strings = list(self)
  162. return (StringStore, (strings,), None, None, None)
  163. def to_disk(self, path):
  164. """Save the current state to a directory.
  165. path (unicode or Path): A path to a directory, which will be created if
  166. it doesn't exist. Paths may be either strings or Path-like objects.
  167. """
  168. path = util.ensure_path(path)
  169. strings = list(self)
  170. with path.open('w') as file_:
  171. file_.write(json_dumps(strings))
  172. def from_disk(self, path):
  173. """Loads state from a directory. Modifies the object in place and
  174. returns it.
  175. path (unicode or Path): A path to a directory. Paths may be either
  176. strings or `Path`-like objects.
  177. RETURNS (StringStore): The modified `StringStore` object.
  178. """
  179. path = util.ensure_path(path)
  180. with path.open('r') as file_:
  181. strings = ujson.load(file_)
  182. prev = list(self)
  183. self._reset_and_load(strings)
  184. for word in prev:
  185. self.add(word)
  186. return self
  187. def to_bytes(self, **exclude):
  188. """Serialize the current state to a binary string.
  189. **exclude: Named attributes to prevent from being serialized.
  190. RETURNS (bytes): The serialized form of the `StringStore` object.
  191. """
  192. return json_dumps(list(self))
  193. def from_bytes(self, bytes_data, **exclude):
  194. """Load state from a binary string.
  195. bytes_data (bytes): The data to load from.
  196. **exclude: Named attributes to prevent from being loaded.
  197. RETURNS (StringStore): The `StringStore` object.
  198. """
  199. strings = ujson.loads(bytes_data)
  200. prev = list(self)
  201. self._reset_and_load(strings)
  202. for word in prev:
  203. self.add(word)
  204. return self
  205. def _reset_and_load(self, strings):
  206. self.mem = Pool()
  207. self._map = PreshMap()
  208. self.keys.clear()
  209. self.hits.clear()
  210. for string in strings:
  211. self.add(string)
  212. def _cleanup_stale_strings(self, excepted):
  213. """
  214. excepted (list): Strings that should not be removed.
  215. RETURNS (keys, strings): Dropped strings and keys that can be dropped from other places
  216. """
  217. if self.hits.size() == 0:
  218. # If we don't have any hits, just skip cleanup
  219. return
  220. cdef vector[hash_t] tmp
  221. dropped_strings = []
  222. dropped_keys = []
  223. for i in range(self.keys.size()):
  224. key = self.keys[i]
  225. # Here we cannot use __getitem__ because it also set hit.
  226. utf8str = <Utf8Str*>self._map.get(key)
  227. value = decode_Utf8Str(utf8str)
  228. if self.hits.count(key) != 0 or value in excepted:
  229. tmp.push_back(key)
  230. else:
  231. dropped_keys.append(key)
  232. dropped_strings.append(value)
  233. self.keys.swap(tmp)
  234. strings = list(self)
  235. self._reset_and_load(strings)
  236. # Here we have strings but hits to it should be reseted
  237. self.hits.clear()
  238. return dropped_keys, dropped_strings
  239. cdef const Utf8Str* intern_unicode(self, unicode py_string):
  240. # 0 means missing, but we don't bother offsetting the index.
  241. cdef bytes byte_string = py_string.encode('utf8')
  242. return self._intern_utf8(byte_string, len(byte_string))
  243. @cython.final
  244. cdef const Utf8Str* _intern_utf8(self, char* utf8_string, int length):
  245. # TODO: This function's API/behaviour is an unholy mess...
  246. # 0 means missing, but we don't bother offsetting the index.
  247. cdef hash_t key = hash_utf8(utf8_string, length)
  248. cdef Utf8Str* value = <Utf8Str*>self._map.get(key)
  249. if value is not NULL:
  250. return value
  251. value = _allocate(self.mem, <unsigned char*>utf8_string, length)
  252. self._map.set(key, value)
  253. self.hits.insert(key)
  254. self.keys.push_back(key)
  255. return value