Improved PriorityQueue and PriorityStack that boost the performance again. Also updated the unittest for those datatypes.

--HG--
extra : convert_revision : svn%3Ab85faa78-f9eb-468e-a121-7cced6da292c%40992

scrapy/trunk/scrapy/tests/test_utils_datatypes.py

@@ -2,41 +2,75 @@ import unittest
 
 from scrapy.utils.datatypes import PriorityQueue, PriorityStack
 
-class DatatypesTestCase(unittest.TestCase):
+# (ITEM, PRIORITY)
+INPUT = [(1, -5), (30, -1), (80, -3), (4, 1), (6, 3), (20, 0), (50, -1)]
 
-    def test_priority_queue(self):
+class PriorityQueueTestCase(unittest.TestCase):
 
-        input = [('five', 5), ('three-1', 3), ('three-2', 3), ('six', 6), ('one', 1)]
-        output = [('one', 1), ('three-1', 3), ('three-2', 3), ('five', 5), ('six', 6)]
+    output = [(1, -5), (80, -3), (30, -1), (50, -1), (20, 0), (4, 1), (6, 3)]
 
+    def test_popping(self):
         pq = PriorityQueue()
-        assert not bool(pq)
-        for item, prio in input:
-            pq.push(item, prio)
-        out = []
-        assert bool(pq)
-        self.assertEqual(len(pq), len(input))
+        for item, pr in INPUT:
+            pq.push(item, pr)
+        l = []
         while pq:
-            out.append(pq.pop())
-        self.assertEqual(len(pq), 0)
-        self.assertEqual(out, output)
+            l.append(pq.pop())
+        self.assertEquals(l, self.output)
 
-    def test_priority_stack(self):
+    def test_iter(self):
+        pq = PriorityQueue()
+        for item, pr in INPUT:
+            pq.push(item, pr)
+        result = [x for x in pq]
+        self.assertEquals(result, self.output) 
+    
+    def test_nonzero(self):
+        pq = PriorityQueue()
+        pq.push(80, -1)
+        pq.push(20, 0)
+        pq.push(30, 1)
 
-        input = [('five', 5), ('three-1', 3), ('three-2', 3), ('six', 6), ('one', 1)]
-        output = [('one', 1), ('three-2', 3), ('three-1', 3), ('five', 5), ('six', 6)]
+        pq.pop()
+        self.assertEquals(bool(pq), True) 
+        pq.pop()
+        self.assertEquals(bool(pq), True) 
+        pq.pop()
+        self.assertEquals(bool(pq), False) 
 
+    def test_len(self):
+        pq = PriorityQueue()
+        pq.push(80, -1)
+        pq.push(20, 0)
+        pq.push(30, 1)
+
+        self.assertEquals(len(pq), 3) 
+        pq.pop()
+        self.assertEquals(len(pq), 2) 
+        pq.pop()
+        self.assertEquals(len(pq), 1) 
+        pq.pop()
+        self.assertEquals(len(pq), 0) 
+
+class PriorityStackTestCase(unittest.TestCase):
+
+    output = [(1, -5), (80, -3), (50, -1), (30, -1), (20, 0), (4, 1), (6, 3)]
+
+    def test_popping(self):
         pq = PriorityStack()
-        assert not bool(pq)
-        for item, prio in input:
-            pq.push(item, prio)
-        assert bool(pq)
-        out = []
-        self.assertEqual(len(pq), len(input))
+        for item, pr in INPUT:
+            pq.push(item, pr)
+        l = []
         while pq:
-            out.append(pq.pop())
-        self.assertEqual(len(pq), 0)
-        self.assertEqual(out, output)
+            l.append(pq.pop())
+        self.assertEquals(l, self.output)
 
+    def test_iter(self):
+        pq = PriorityStack()
+        for item, pr in INPUT:
+            pq.push(item, pr)
+        result = [x for x in pq]
+        self.assertEquals(result, self.output) 
+    
 if __name__ == "__main__":
     unittest.main()

scrapy/trunk/scrapy/utils/datatypes.py

@@ -7,6 +7,7 @@ This module must not depend on any module outside the Standard Library.
 
 import copy
 from collections import deque, defaultdict
+from itertools import chain
 
 class MultiValueDictKeyError(KeyError):
     pass
@@ -229,52 +230,75 @@ class CaselessDict(dict):
         return dict.pop(self, self.normkey(key), *args)
 
 
-
 class PriorityQueue(object):
-    """
-    Faster priority queue implementation, based on a dictionary of lists
-    @author: Federico Feroldi <federico@cloudify.me>
-    """
+    """Priority queue using a deque for priority 0"""
+
     def __init__(self):
-        self.items = defaultdict(deque)
+        self.negitems = defaultdict(deque)
+        self.pzero = deque()
+        self.positems = defaultdict(deque)
 
     def push(self, item, priority=0):
-        self.items[priority].appendleft(item)
+        if priority == 0:
+            self.pzero.appendleft(item)
+        elif priority < 0:
+            self.negitems[priority].appendleft(item)
+        else:
+            self.positems[priority].appendleft(item)
 
     def pop(self):
-        priorities = self.items.keys()
-        priorities.sort()
-        
-        for priority in priorities:
-            if len(self.items[priority]) > 0:
-                return (self.items[priority].pop(), priority)
-            
-        raise IndexError
+        if self.negitems:
+            priorities = self.negitems.keys()
+            priorities.sort()
+            for priority in priorities:
+                deq = self.negitems[priority]
+                if deq:
+                    t = (deq.pop(), priority)
+                    if not deq:
+                        del self.negitems[priority]
+                    return t
+        elif self.pzero:
+            return (self.pzero.pop(), 0)
+        else:
+            priorities = self.positems.keys()
+            priorities.sort()
+            for priority in priorities:
+                deq = self.positems[priority]
+                if deq:
+                    t = (deq.pop(), priority)
+                    if not deq:
+                        del self.positems[priority]
+                    return t
+        raise IndexError("pop from an empty queue")
 
     def __len__(self):
-        totlen = 0
-        for q in self.items.values():
-            totlen += len(q)
-        return totlen
+        total = sum(len(v) for v in self.negitems.values()) + \
+                len(self.pzero) + \
+                sum(len(v) for v in self.positems.values())
+        return total
 
     def __iter__(self):
-        priorities = self.items.keys()
-        priorities.sort()
-        
-        for priority in priorities:
-            for i in self.items[priority]:
-                yield (i, priority)
-            
+        gen_negs = ((i, priority) 
+                    for priority in sorted(self.negitems.keys())
+                    for i in reversed(self.negitems[priority]))
+        gen_zeros = ((item,0) for item in self.pzero)
+        gen_pos = ((i, priority) 
+                    for priority in sorted(self.positems.keys())
+                    for i in reversed(self.positems[priority]))
+        return chain(gen_negs, gen_zeros, gen_pos)
+
     def __nonzero__(self):
-        for q in self.items.values():
-            if len(q) > 0:
-                return True
-        return False
+        return bool(self.negitems or self.pzero or self.positems)
 
 class PriorityStack(PriorityQueue):
     """A simple priority stack which is similar to PriorityQueue but pops its
     items in reverse order (for the same priority)"""
 
     def push(self, item, priority=0):
-        self.items[priority].append(item)
+        if priority == 0:
+            self.pzero.append(item)
+        elif priority < 0:
+            self.negitems[priority].append(item)
+        else:
+            self.positems[priority].append(item)