collections.py 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661
  1. # Natural Language Toolkit: Collections
  2. #
  3. # Copyright (C) 2001-2020 NLTK Project
  4. # Author: Steven Bird <stevenbird1@gmail.com>
  5. # URL: <http://nltk.org/>
  6. # For license information, see LICENSE.TXT
  7. import bisect
  8. from itertools import islice, chain
  9. from functools import total_ordering
  10. # this unused import is for python 2.7
  11. from collections import defaultdict, deque, Counter
  12. from nltk.internals import slice_bounds, raise_unorderable_types
  13. ##########################################################################
  14. # Ordered Dictionary
  15. ##########################################################################
  16. class OrderedDict(dict):
  17. def __init__(self, data=None, **kwargs):
  18. self._keys = self.keys(data, kwargs.get("keys"))
  19. self._default_factory = kwargs.get("default_factory")
  20. if data is None:
  21. dict.__init__(self)
  22. else:
  23. dict.__init__(self, data)
  24. def __delitem__(self, key):
  25. dict.__delitem__(self, key)
  26. self._keys.remove(key)
  27. def __getitem__(self, key):
  28. try:
  29. return dict.__getitem__(self, key)
  30. except KeyError:
  31. return self.__missing__(key)
  32. def __iter__(self):
  33. return (key for key in self.keys())
  34. def __missing__(self, key):
  35. if not self._default_factory and key not in self._keys:
  36. raise KeyError()
  37. return self._default_factory()
  38. def __setitem__(self, key, item):
  39. dict.__setitem__(self, key, item)
  40. if key not in self._keys:
  41. self._keys.append(key)
  42. def clear(self):
  43. dict.clear(self)
  44. self._keys.clear()
  45. def copy(self):
  46. d = dict.copy(self)
  47. d._keys = self._keys
  48. return d
  49. def items(self):
  50. # returns iterator under python 3 and list under python 2
  51. return zip(self.keys(), self.values())
  52. def keys(self, data=None, keys=None):
  53. if data:
  54. if keys:
  55. assert isinstance(keys, list)
  56. assert len(data) == len(keys)
  57. return keys
  58. else:
  59. assert (
  60. isinstance(data, dict)
  61. or isinstance(data, OrderedDict)
  62. or isinstance(data, list)
  63. )
  64. if isinstance(data, dict) or isinstance(data, OrderedDict):
  65. return data.keys()
  66. elif isinstance(data, list):
  67. return [key for (key, value) in data]
  68. elif "_keys" in self.__dict__:
  69. return self._keys
  70. else:
  71. return []
  72. def popitem(self):
  73. if not self._keys:
  74. raise KeyError()
  75. key = self._keys.pop()
  76. value = self[key]
  77. del self[key]
  78. return (key, value)
  79. def setdefault(self, key, failobj=None):
  80. dict.setdefault(self, key, failobj)
  81. if key not in self._keys:
  82. self._keys.append(key)
  83. def update(self, data):
  84. dict.update(self, data)
  85. for key in self.keys(data):
  86. if key not in self._keys:
  87. self._keys.append(key)
  88. def values(self):
  89. # returns iterator under python 3
  90. return map(self.get, self._keys)
  91. ######################################################################
  92. # Lazy Sequences
  93. ######################################################################
  94. @total_ordering
  95. class AbstractLazySequence(object):
  96. """
  97. An abstract base class for read-only sequences whose values are
  98. computed as needed. Lazy sequences act like tuples -- they can be
  99. indexed, sliced, and iterated over; but they may not be modified.
  100. The most common application of lazy sequences in NLTK is for
  101. corpus view objects, which provide access to the contents of a
  102. corpus without loading the entire corpus into memory, by loading
  103. pieces of the corpus from disk as needed.
  104. The result of modifying a mutable element of a lazy sequence is
  105. undefined. In particular, the modifications made to the element
  106. may or may not persist, depending on whether and when the lazy
  107. sequence caches that element's value or reconstructs it from
  108. scratch.
  109. Subclasses are required to define two methods: ``__len__()``
  110. and ``iterate_from()``.
  111. """
  112. def __len__(self):
  113. """
  114. Return the number of tokens in the corpus file underlying this
  115. corpus view.
  116. """
  117. raise NotImplementedError("should be implemented by subclass")
  118. def iterate_from(self, start):
  119. """
  120. Return an iterator that generates the tokens in the corpus
  121. file underlying this corpus view, starting at the token number
  122. ``start``. If ``start>=len(self)``, then this iterator will
  123. generate no tokens.
  124. """
  125. raise NotImplementedError("should be implemented by subclass")
  126. def __getitem__(self, i):
  127. """
  128. Return the *i* th token in the corpus file underlying this
  129. corpus view. Negative indices and spans are both supported.
  130. """
  131. if isinstance(i, slice):
  132. start, stop = slice_bounds(self, i)
  133. return LazySubsequence(self, start, stop)
  134. else:
  135. # Handle negative indices
  136. if i < 0:
  137. i += len(self)
  138. if i < 0:
  139. raise IndexError("index out of range")
  140. # Use iterate_from to extract it.
  141. try:
  142. return next(self.iterate_from(i))
  143. except StopIteration:
  144. raise IndexError("index out of range")
  145. def __iter__(self):
  146. """Return an iterator that generates the tokens in the corpus
  147. file underlying this corpus view."""
  148. return self.iterate_from(0)
  149. def count(self, value):
  150. """Return the number of times this list contains ``value``."""
  151. return sum(1 for elt in self if elt == value)
  152. def index(self, value, start=None, stop=None):
  153. """Return the index of the first occurrence of ``value`` in this
  154. list that is greater than or equal to ``start`` and less than
  155. ``stop``. Negative start and stop values are treated like negative
  156. slice bounds -- i.e., they count from the end of the list."""
  157. start, stop = slice_bounds(self, slice(start, stop))
  158. for i, elt in enumerate(islice(self, start, stop)):
  159. if elt == value:
  160. return i + start
  161. raise ValueError("index(x): x not in list")
  162. def __contains__(self, value):
  163. """Return true if this list contains ``value``."""
  164. return bool(self.count(value))
  165. def __add__(self, other):
  166. """Return a list concatenating self with other."""
  167. return LazyConcatenation([self, other])
  168. def __radd__(self, other):
  169. """Return a list concatenating other with self."""
  170. return LazyConcatenation([other, self])
  171. def __mul__(self, count):
  172. """Return a list concatenating self with itself ``count`` times."""
  173. return LazyConcatenation([self] * count)
  174. def __rmul__(self, count):
  175. """Return a list concatenating self with itself ``count`` times."""
  176. return LazyConcatenation([self] * count)
  177. _MAX_REPR_SIZE = 60
  178. def __repr__(self):
  179. """
  180. Return a string representation for this corpus view that is
  181. similar to a list's representation; but if it would be more
  182. than 60 characters long, it is truncated.
  183. """
  184. pieces = []
  185. length = 5
  186. for elt in self:
  187. pieces.append(repr(elt))
  188. length += len(pieces[-1]) + 2
  189. if length > self._MAX_REPR_SIZE and len(pieces) > 2:
  190. return "[%s, ...]" % ", ".join(pieces[:-1])
  191. return "[%s]" % ", ".join(pieces)
  192. def __eq__(self, other):
  193. return type(self) == type(other) and list(self) == list(other)
  194. def __ne__(self, other):
  195. return not self == other
  196. def __lt__(self, other):
  197. if type(other) != type(self):
  198. raise_unorderable_types("<", self, other)
  199. return list(self) < list(other)
  200. def __hash__(self):
  201. """
  202. :raise ValueError: Corpus view objects are unhashable.
  203. """
  204. raise ValueError("%s objects are unhashable" % self.__class__.__name__)
  205. class LazySubsequence(AbstractLazySequence):
  206. """
  207. A subsequence produced by slicing a lazy sequence. This slice
  208. keeps a reference to its source sequence, and generates its values
  209. by looking them up in the source sequence.
  210. """
  211. MIN_SIZE = 100
  212. """
  213. The minimum size for which lazy slices should be created. If
  214. ``LazySubsequence()`` is called with a subsequence that is
  215. shorter than ``MIN_SIZE``, then a tuple will be returned instead.
  216. """
  217. def __new__(cls, source, start, stop):
  218. """
  219. Construct a new slice from a given underlying sequence. The
  220. ``start`` and ``stop`` indices should be absolute indices --
  221. i.e., they should not be negative (for indexing from the back
  222. of a list) or greater than the length of ``source``.
  223. """
  224. # If the slice is small enough, just use a tuple.
  225. if stop - start < cls.MIN_SIZE:
  226. return list(islice(source.iterate_from(start), stop - start))
  227. else:
  228. return object.__new__(cls)
  229. def __init__(self, source, start, stop):
  230. self._source = source
  231. self._start = start
  232. self._stop = stop
  233. def __len__(self):
  234. return self._stop - self._start
  235. def iterate_from(self, start):
  236. return islice(
  237. self._source.iterate_from(start + self._start), max(0, len(self) - start)
  238. )
  239. class LazyConcatenation(AbstractLazySequence):
  240. """
  241. A lazy sequence formed by concatenating a list of lists. This
  242. underlying list of lists may itself be lazy. ``LazyConcatenation``
  243. maintains an index that it uses to keep track of the relationship
  244. between offsets in the concatenated lists and offsets in the
  245. sublists.
  246. """
  247. def __init__(self, list_of_lists):
  248. self._list = list_of_lists
  249. self._offsets = [0]
  250. def __len__(self):
  251. if len(self._offsets) <= len(self._list):
  252. for tok in self.iterate_from(self._offsets[-1]):
  253. pass
  254. return self._offsets[-1]
  255. def iterate_from(self, start_index):
  256. if start_index < self._offsets[-1]:
  257. sublist_index = bisect.bisect_right(self._offsets, start_index) - 1
  258. else:
  259. sublist_index = len(self._offsets) - 1
  260. index = self._offsets[sublist_index]
  261. # Construct an iterator over the sublists.
  262. if isinstance(self._list, AbstractLazySequence):
  263. sublist_iter = self._list.iterate_from(sublist_index)
  264. else:
  265. sublist_iter = islice(self._list, sublist_index, None)
  266. for sublist in sublist_iter:
  267. if sublist_index == (len(self._offsets) - 1):
  268. assert (
  269. index + len(sublist) >= self._offsets[-1]
  270. ), "offests not monotonic increasing!"
  271. self._offsets.append(index + len(sublist))
  272. else:
  273. assert self._offsets[sublist_index + 1] == index + len(
  274. sublist
  275. ), "inconsistent list value (num elts)"
  276. for value in sublist[max(0, start_index - index) :]:
  277. yield value
  278. index += len(sublist)
  279. sublist_index += 1
  280. class LazyMap(AbstractLazySequence):
  281. """
  282. A lazy sequence whose elements are formed by applying a given
  283. function to each element in one or more underlying lists. The
  284. function is applied lazily -- i.e., when you read a value from the
  285. list, ``LazyMap`` will calculate that value by applying its
  286. function to the underlying lists' value(s). ``LazyMap`` is
  287. essentially a lazy version of the Python primitive function
  288. ``map``. In particular, the following two expressions are
  289. equivalent:
  290. >>> from nltk.collections import LazyMap
  291. >>> function = str
  292. >>> sequence = [1,2,3]
  293. >>> map(function, sequence) # doctest: +SKIP
  294. ['1', '2', '3']
  295. >>> list(LazyMap(function, sequence))
  296. ['1', '2', '3']
  297. Like the Python ``map`` primitive, if the source lists do not have
  298. equal size, then the value None will be supplied for the
  299. 'missing' elements.
  300. Lazy maps can be useful for conserving memory, in cases where
  301. individual values take up a lot of space. This is especially true
  302. if the underlying list's values are constructed lazily, as is the
  303. case with many corpus readers.
  304. A typical example of a use case for this class is performing
  305. feature detection on the tokens in a corpus. Since featuresets
  306. are encoded as dictionaries, which can take up a lot of memory,
  307. using a ``LazyMap`` can significantly reduce memory usage when
  308. training and running classifiers.
  309. """
  310. def __init__(self, function, *lists, **config):
  311. """
  312. :param function: The function that should be applied to
  313. elements of ``lists``. It should take as many arguments
  314. as there are ``lists``.
  315. :param lists: The underlying lists.
  316. :param cache_size: Determines the size of the cache used
  317. by this lazy map. (default=5)
  318. """
  319. if not lists:
  320. raise TypeError("LazyMap requires at least two args")
  321. self._lists = lists
  322. self._func = function
  323. self._cache_size = config.get("cache_size", 5)
  324. self._cache = {} if self._cache_size > 0 else None
  325. # If you just take bool() of sum() here _all_lazy will be true just
  326. # in case n >= 1 list is an AbstractLazySequence. Presumably this
  327. # isn't what's intended.
  328. self._all_lazy = sum(
  329. isinstance(lst, AbstractLazySequence) for lst in lists
  330. ) == len(lists)
  331. def iterate_from(self, index):
  332. # Special case: one lazy sublist
  333. if len(self._lists) == 1 and self._all_lazy:
  334. for value in self._lists[0].iterate_from(index):
  335. yield self._func(value)
  336. return
  337. # Special case: one non-lazy sublist
  338. elif len(self._lists) == 1:
  339. while True:
  340. try:
  341. yield self._func(self._lists[0][index])
  342. except IndexError:
  343. return
  344. index += 1
  345. # Special case: n lazy sublists
  346. elif self._all_lazy:
  347. iterators = [lst.iterate_from(index) for lst in self._lists]
  348. while True:
  349. elements = []
  350. for iterator in iterators:
  351. try:
  352. elements.append(next(iterator))
  353. except: # FIXME: What is this except really catching? StopIteration?
  354. elements.append(None)
  355. if elements == [None] * len(self._lists):
  356. return
  357. yield self._func(*elements)
  358. index += 1
  359. # general case
  360. else:
  361. while True:
  362. try:
  363. elements = [lst[index] for lst in self._lists]
  364. except IndexError:
  365. elements = [None] * len(self._lists)
  366. for i, lst in enumerate(self._lists):
  367. try:
  368. elements[i] = lst[index]
  369. except IndexError:
  370. pass
  371. if elements == [None] * len(self._lists):
  372. return
  373. yield self._func(*elements)
  374. index += 1
  375. def __getitem__(self, index):
  376. if isinstance(index, slice):
  377. sliced_lists = [lst[index] for lst in self._lists]
  378. return LazyMap(self._func, *sliced_lists)
  379. else:
  380. # Handle negative indices
  381. if index < 0:
  382. index += len(self)
  383. if index < 0:
  384. raise IndexError("index out of range")
  385. # Check the cache
  386. if self._cache is not None and index in self._cache:
  387. return self._cache[index]
  388. # Calculate the value
  389. try:
  390. val = next(self.iterate_from(index))
  391. except StopIteration:
  392. raise IndexError("index out of range")
  393. # Update the cache
  394. if self._cache is not None:
  395. if len(self._cache) > self._cache_size:
  396. self._cache.popitem() # discard random entry
  397. self._cache[index] = val
  398. # Return the value
  399. return val
  400. def __len__(self):
  401. return max(len(lst) for lst in self._lists)
  402. class LazyZip(LazyMap):
  403. """
  404. A lazy sequence whose elements are tuples, each containing the i-th
  405. element from each of the argument sequences. The returned list is
  406. truncated in length to the length of the shortest argument sequence. The
  407. tuples are constructed lazily -- i.e., when you read a value from the
  408. list, ``LazyZip`` will calculate that value by forming a tuple from
  409. the i-th element of each of the argument sequences.
  410. ``LazyZip`` is essentially a lazy version of the Python primitive function
  411. ``zip``. In particular, an evaluated LazyZip is equivalent to a zip:
  412. >>> from nltk.collections import LazyZip
  413. >>> sequence1, sequence2 = [1, 2, 3], ['a', 'b', 'c']
  414. >>> zip(sequence1, sequence2) # doctest: +SKIP
  415. [(1, 'a'), (2, 'b'), (3, 'c')]
  416. >>> list(LazyZip(sequence1, sequence2))
  417. [(1, 'a'), (2, 'b'), (3, 'c')]
  418. >>> sequences = [sequence1, sequence2, [6,7,8,9]]
  419. >>> list(zip(*sequences)) == list(LazyZip(*sequences))
  420. True
  421. Lazy zips can be useful for conserving memory in cases where the argument
  422. sequences are particularly long.
  423. A typical example of a use case for this class is combining long sequences
  424. of gold standard and predicted values in a classification or tagging task
  425. in order to calculate accuracy. By constructing tuples lazily and
  426. avoiding the creation of an additional long sequence, memory usage can be
  427. significantly reduced.
  428. """
  429. def __init__(self, *lists):
  430. """
  431. :param lists: the underlying lists
  432. :type lists: list(list)
  433. """
  434. LazyMap.__init__(self, lambda *elts: elts, *lists)
  435. def iterate_from(self, index):
  436. iterator = LazyMap.iterate_from(self, index)
  437. while index < len(self):
  438. yield next(iterator)
  439. index += 1
  440. return
  441. def __len__(self):
  442. return min(len(lst) for lst in self._lists)
  443. class LazyEnumerate(LazyZip):
  444. """
  445. A lazy sequence whose elements are tuples, each ontaining a count (from
  446. zero) and a value yielded by underlying sequence. ``LazyEnumerate`` is
  447. useful for obtaining an indexed list. The tuples are constructed lazily
  448. -- i.e., when you read a value from the list, ``LazyEnumerate`` will
  449. calculate that value by forming a tuple from the count of the i-th
  450. element and the i-th element of the underlying sequence.
  451. ``LazyEnumerate`` is essentially a lazy version of the Python primitive
  452. function ``enumerate``. In particular, the following two expressions are
  453. equivalent:
  454. >>> from nltk.collections import LazyEnumerate
  455. >>> sequence = ['first', 'second', 'third']
  456. >>> list(enumerate(sequence))
  457. [(0, 'first'), (1, 'second'), (2, 'third')]
  458. >>> list(LazyEnumerate(sequence))
  459. [(0, 'first'), (1, 'second'), (2, 'third')]
  460. Lazy enumerations can be useful for conserving memory in cases where the
  461. argument sequences are particularly long.
  462. A typical example of a use case for this class is obtaining an indexed
  463. list for a long sequence of values. By constructing tuples lazily and
  464. avoiding the creation of an additional long sequence, memory usage can be
  465. significantly reduced.
  466. """
  467. def __init__(self, lst):
  468. """
  469. :param lst: the underlying list
  470. :type lst: list
  471. """
  472. LazyZip.__init__(self, range(len(lst)), lst)
  473. class LazyIteratorList(AbstractLazySequence):
  474. """
  475. Wraps an iterator, loading its elements on demand
  476. and making them subscriptable.
  477. __repr__ displays only the first few elements.
  478. """
  479. def __init__(self, it, known_len=None):
  480. self._it = it
  481. self._len = known_len
  482. self._cache = []
  483. def __len__(self):
  484. if self._len:
  485. return self._len
  486. for x in self.iterate_from(len(self._cache)):
  487. pass
  488. self._len = len(self._cache)
  489. return self._len
  490. def iterate_from(self, start):
  491. """Create a new iterator over this list starting at the given offset."""
  492. while len(self._cache) < start:
  493. v = next(self._it)
  494. self._cache.append(v)
  495. i = start
  496. while i < len(self._cache):
  497. yield self._cache[i]
  498. i += 1
  499. while True:
  500. v = next(self._it)
  501. self._cache.append(v)
  502. yield v
  503. i += 1
  504. def __add__(self, other):
  505. """Return a list concatenating self with other."""
  506. return type(self)(chain(self, other))
  507. def __radd__(self, other):
  508. """Return a list concatenating other with self."""
  509. return type(self)(chain(other, self))
  510. ######################################################################
  511. # Trie Implementation
  512. ######################################################################
  513. class Trie(dict):
  514. """A Trie implementation for strings"""
  515. LEAF = True
  516. def __init__(self, strings=None):
  517. """Builds a Trie object, which is built around a ``dict``
  518. If ``strings`` is provided, it will add the ``strings``, which
  519. consist of a ``list`` of ``strings``, to the Trie.
  520. Otherwise, it'll construct an empty Trie.
  521. :param strings: List of strings to insert into the trie
  522. (Default is ``None``)
  523. :type strings: list(str)
  524. """
  525. super(Trie, self).__init__()
  526. if strings:
  527. for string in strings:
  528. self.insert(string)
  529. def insert(self, string):
  530. """Inserts ``string`` into the Trie
  531. :param string: String to insert into the trie
  532. :type string: str
  533. :Example:
  534. >>> from nltk.collections import Trie
  535. >>> trie = Trie(["abc", "def"])
  536. >>> expected = {'a': {'b': {'c': {True: None}}}, \
  537. 'd': {'e': {'f': {True: None}}}}
  538. >>> trie == expected
  539. True
  540. """
  541. if len(string):
  542. self[string[0]].insert(string[1:])
  543. else:
  544. # mark the string is complete
  545. self[Trie.LEAF] = None
  546. def __missing__(self, key):
  547. self[key] = Trie()
  548. return self[key]