diff -r de35eae9048a Lib/queue.py --- a/Lib/queue.py Wed Apr 24 23:33:20 2013 +0200 +++ b/Lib/queue.py Thu Apr 25 20:58:20 2013 +0100 @@ -6,6 +6,7 @@ import dummy_threading as threading from collections import deque from heapq import heappush, heappop +from itertools import count try: from time import monotonic as time except ImportError: @@ -220,17 +221,21 @@ Entries are typically tuples of the form: (priority number, data). ''' - def _init(self, maxsize): + def _init(self, maxsize, key=None): self.queue = [] + if key is None: + key = lambda x: x + self.key = key + self.counter = count() def _qsize(self): return len(self.queue) def _put(self, item): - heappush(self.queue, item) + heappush(self.queue, (self.key(item), next(self.counter), item)) def _get(self): - return heappop(self.queue) + return heappop(self.queue)[2] class LifoQueue(Queue): diff -r de35eae9048a Lib/test/test_queue.py --- a/Lib/test/test_queue.py Wed Apr 24 23:33:20 2013 +0200 +++ b/Lib/test/test_queue.py Thu Apr 25 20:58:20 2013 +0100 @@ -354,6 +354,37 @@ self.failing_queue_test(q) +class TestKeyedPriorityQueue(unittest.TestCase): + def test_key(self): + q = PriorityQueue(key=len) + q.put([1, 2, 3, 4]) + q.put([5, 6]) + q.put([7, 8, 9]) + self.assertEqual(q.get(), [5, 6]) + self.assertEqual(q.get(), [7, 8, 9]) + self.assertEqual(q.get(), [1, 2, 3, 4]) + + def test_lifo_within_equal_priority(self): + test_items = [10, 19, 14, 12, 17, 15, 16, 5, 4, + 2, 18, 11, 13, 0, 1, 3, 7, 6, 8, 9] + expected_order = [5, 4, 2, 0, 1, 3, 7, 6, 8, 9, + 10, 19, 14, 12, 17, 15, 16, 18, 11, 13] + q = PriorityQueue(key=lambda n: n // 10) + for item in test_items: + q.put([item]) + for item in expected_order: + self.assertEqual(q.get(), item) + + def test_unorderable_items(self): + q = PriorityQueue(key=abs) + q.put(3 + 4j) + q.put(1j) + q.put(5 + 12j) + self.assertEqual(q.get(), 1j) + self.assertEqual(q.get(), 3 + 4j) + self.assertEqual(q.get(), 5 + 12j) + + def test_main(): support.run_unittest(QueueTest, LifoQueueTest, PriorityQueueTest, FailingQueueTest)