queueutils.py 7.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225
  1. # -*- coding: utf-8 -*-
  2. # Copyright (c) 2013, Mahmoud Hashemi
  3. #
  4. # Redistribution and use in source and binary forms, with or without
  5. # modification, are permitted provided that the following conditions are
  6. # met:
  7. #
  8. # * Redistributions of source code must retain the above copyright
  9. # notice, this list of conditions and the following disclaimer.
  10. #
  11. # * Redistributions in binary form must reproduce the above
  12. # copyright notice, this list of conditions and the following
  13. # disclaimer in the documentation and/or other materials provided
  14. # with the distribution.
  15. #
  16. # * The names of the contributors may not be used to endorse or
  17. # promote products derived from this software without specific
  18. # prior written permission.
  19. #
  20. # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  21. # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  22. # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  23. # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  24. # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  25. # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  26. # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  27. # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  28. # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  29. # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  30. # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  31. """Python comes with a many great data structures, from :class:`dict`
  32. to :class:`collections.deque`, and no shortage of serviceable
  33. algorithm implementations, from :func:`sorted` to :mod:`bisect`. But
  34. priority queues are curiously relegated to an example documented in
  35. :mod:`heapq`. Even there, the approach presented is not full-featured
  36. and object-oriented. There is a built-in priority queue,
  37. :class:`Queue.PriorityQueue`, but in addition to its austere API, it
  38. carries the double-edged sword of threadsafety, making it fine for
  39. multi-threaded, multi-consumer applications, but high-overhead for
  40. cooperative/single-threaded use cases.
  41. The ``queueutils`` module currently provides two Queue
  42. implementations: :class:`HeapPriorityQueue`, based on a heap, and
  43. :class:`SortedPriorityQueue`, based on a sorted list. Both use a
  44. unified API based on :class:`BasePriorityQueue` to facilitate testing
  45. the slightly different performance characteristics on various
  46. application use cases.
  47. >>> pq = PriorityQueue()
  48. >>> pq.add('low priority task', 0)
  49. >>> pq.add('high priority task', 2)
  50. >>> pq.add('medium priority task 1', 1)
  51. >>> pq.add('medium priority task 2', 1)
  52. >>> len(pq)
  53. 4
  54. >>> pq.pop()
  55. 'high priority task'
  56. >>> pq.peek()
  57. 'medium priority task 1'
  58. >>> len(pq)
  59. 3
  60. """
  61. from heapq import heappush, heappop
  62. from bisect import insort
  63. import itertools
  64. try:
  65. from .typeutils import make_sentinel
  66. _REMOVED = make_sentinel(var_name='_REMOVED')
  67. except ImportError:
  68. _REMOVED = object()
  69. try:
  70. from .listutils import BList
  71. # see BarrelList docstring for notes
  72. except ImportError:
  73. BList = list
  74. __all__ = ['PriorityQueue', 'BasePriorityQueue',
  75. 'HeapPriorityQueue', 'SortedPriorityQueue']
  76. # TODO: make Base a real abstract class
  77. # TODO: add uniqueification
  78. class BasePriorityQueue(object):
  79. """The abstract base class for the other PriorityQueues in this
  80. module. Override the ``_backend_type`` class attribute, as well as
  81. the :meth:`_push_entry` and :meth:`_pop_entry` staticmethods for
  82. custom subclass behavior. (Don't forget to use
  83. :func:`staticmethod`).
  84. Args:
  85. priority_key (callable): A function that takes *priority* as
  86. passed in by :meth:`add` and returns a real number
  87. representing the effective priority.
  88. """
  89. # negating priority means larger numbers = higher priority
  90. _default_priority_key = staticmethod(lambda p: -float(p or 0))
  91. _backend_type = list
  92. def __init__(self, **kw):
  93. self._pq = self._backend_type()
  94. self._entry_map = {}
  95. self._counter = itertools.count()
  96. self._get_priority = kw.pop('priority_key', self._default_priority_key)
  97. if kw:
  98. raise TypeError('unexpected keyword arguments: %r' % kw.keys())
  99. @staticmethod
  100. def _push_entry(backend, entry):
  101. pass # abstract
  102. @staticmethod
  103. def _pop_entry(backend):
  104. pass # abstract
  105. def add(self, task, priority=None):
  106. """
  107. Add a task to the queue, or change the *task*'s priority if *task*
  108. is already in the queue. *task* can be any hashable object,
  109. and *priority* defaults to ``0``. Higher values representing
  110. higher priority, but this behavior can be controlled by
  111. setting *priority_key* in the constructor.
  112. """
  113. priority = self._get_priority(priority)
  114. if task in self._entry_map:
  115. self.remove(task)
  116. count = next(self._counter)
  117. entry = [priority, count, task]
  118. self._entry_map[task] = entry
  119. self._push_entry(self._pq, entry)
  120. def remove(self, task):
  121. """Remove a task from the priority queue. Raises :exc:`KeyError` if
  122. the *task* is absent.
  123. """
  124. entry = self._entry_map.pop(task)
  125. entry[-1] = _REMOVED
  126. def _cull(self, raise_exc=True):
  127. "Remove entries marked as removed by previous :meth:`remove` calls."
  128. while self._pq:
  129. priority, count, task = self._pq[0]
  130. if task is _REMOVED:
  131. self._pop_entry(self._pq)
  132. continue
  133. return
  134. if raise_exc:
  135. raise IndexError('empty priority queue')
  136. def peek(self, default=_REMOVED):
  137. """Read the next value in the queue without removing it. Returns
  138. *default* on an empty queue, or raises :exc:`KeyError` if
  139. *default* is not set.
  140. """
  141. try:
  142. self._cull()
  143. _, _, task = self._pq[0]
  144. except IndexError:
  145. if default is not _REMOVED:
  146. return default
  147. raise IndexError('peek on empty queue')
  148. return task
  149. def pop(self, default=_REMOVED):
  150. """Remove and return the next value in the queue. Returns *default* on
  151. an empty queue, or raises :exc:`KeyError` if *default* is not
  152. set.
  153. """
  154. try:
  155. self._cull()
  156. _, _, task = self._pop_entry(self._pq)
  157. del self._entry_map[task]
  158. except IndexError:
  159. if default is not _REMOVED:
  160. return default
  161. raise IndexError('pop on empty queue')
  162. return task
  163. def __len__(self):
  164. "Return the number of tasks in the queue."
  165. return len(self._entry_map)
  166. class HeapPriorityQueue(BasePriorityQueue):
  167. """A priority queue inherited from :class:`BasePriorityQueue`,
  168. backed by a list and based on the :func:`heapq.heappop` and
  169. :func:`heapq.heappush` functions in the built-in :mod:`heapq`
  170. module.
  171. """
  172. @staticmethod
  173. def _pop_entry(backend):
  174. return heappop(backend)
  175. @staticmethod
  176. def _push_entry(backend, entry):
  177. heappush(backend, entry)
  178. class SortedPriorityQueue(BasePriorityQueue):
  179. """A priority queue inherited from :class:`BasePriorityQueue`, based
  180. on the :func:`bisect.insort` approach for in-order insertion into
  181. a sorted list.
  182. """
  183. _backend_type = BList
  184. @staticmethod
  185. def _pop_entry(backend):
  186. return backend.pop(0)
  187. @staticmethod
  188. def _push_entry(backend, entry):
  189. insort(backend, entry)
  190. PriorityQueue = SortedPriorityQueue