hacky refactor to have a BasePriorityQueue to make room for SortedPriorityQueue with peek_n, etc.

boltons/queueutils.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 from heapq import heappush, heappop
+from bisect import insort
 import itertools
 
 try:
@@ -9,40 +10,61 @@ try:
 except ImportError:
     _REMOVED = object()
 
-
-__all__ = ['PriorityQueue', 'HeapPriorityQueue']
+try:
+    from listutils import BList
+    list_type = BList  # see BarrelList docstring for notes
+except ImportError:
+    list_type = list
 
 
-class HeapPriorityQueue(object):
-    """
-    Real quick type based on the heapq docs.
-    """
-    def __init__(self):
-        self._pq = []
+__all__ = ['PriorityQueue', 'HeapPriorityQueue', 'SortedPriorityQueue']
+
+
+# TODO: make Base a real abstract class
+# TODO: add uniqueification
+
+
+class BasePriorityQueue(object):
+    # negating priority means larger numbers = higher priority
+    _default_priority_key = staticmethod(lambda p: -int(p or 0))
+    _list_type = list
+
+    def __init__(self, **kw):
+        self._pq = self._list_type()
         self._entry_map = {}
-        self.counter = itertools.count()
+        self._counter = itertools.count()
+        self._get_priority = kw.pop('priority_key', self._default_priority_key)
+
+    @staticmethod
+    def _push_entry(backend, entry):
+        pass  # abstract
+
+    @staticmethod
+    def _pop_entry(backend):
+        pass  # abstract
 
     def add(self, task, priority=None):
-        # larger numbers = higher priority
-        priority = -int(priority or 0)
+        priority = self._get_priority(priority)
         if task in self._entry_map:
-            self.remove_task(task)
-        count = next(self.counter)
+            self.remove(task)
+        count = next(self._counter)
         entry = [priority, count, task]
         self._entry_map[task] = entry
-        heappush(self._pq, entry)
+        self._push_entry(self._pq, entry)
 
     def remove(self, task):
         entry = self._entry_map.pop(task)
         entry[-1] = _REMOVED
 
-    def _cull(self):
+    def _cull(self, raise_exc=True):
         while self._pq:
             priority, count, task = self._pq[0]
             if task is _REMOVED:
-                heappop(self._pq)
+                self._pop_entry(self._pq)
                 continue
             return
+        if raise_exc:
+            raise IndexError('empty priority queue')
 
     def peek(self, default=_REMOVED):
         try:
@@ -57,7 +79,7 @@ class HeapPriorityQueue(object):
     def pop(self, default=_REMOVED):
         try:
             self._cull()
-            _, _, task = heappop(self._pq)
+            _, _, task = self._pop_entry(self._pq)
             del self._entry_map[task]
         except IndexError:
             if default is not _REMOVED:
@@ -69,4 +91,38 @@ class HeapPriorityQueue(object):
         return len(self._entry_map)
 
 
+class HeapPriorityQueue(BasePriorityQueue):
+    @staticmethod
+    def _pop_entry(backend):
+        return heappop(backend)
+
+    @staticmethod
+    def _push_entry(backend, entry):
+        heappush(backend, entry)
+
+
+class SortedPriorityQueue(BasePriorityQueue):
+    pass
+
+
 PriorityQueue = HeapPriorityQueue
+
+# tests
+
+def main():
+    pq = PriorityQueue()
+    func = lambda x: x
+    pq.add(func)
+    pq.remove(func)
+    pq.add(func)
+    pq.add(func)
+    assert len(pq) == 1
+    assert func == pq.pop()
+    assert len(pq) == 0
+
+if __name__ == '__main__':
+    try:
+        main()
+    except Exception as e:
+        import pdb;pdb.post_mortem()
+        raise