diff --git a/scrapy/trunk/scrapy/tests/test_utils_datatypes.py b/scrapy/trunk/scrapy/tests/test_utils_datatypes.py index babf17697..8cce73f53 100644 --- a/scrapy/trunk/scrapy/tests/test_utils_datatypes.py +++ b/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() diff --git a/scrapy/trunk/scrapy/utils/datatypes.py b/scrapy/trunk/scrapy/utils/datatypes.py index 8a2997656..2f42ba739 100644 --- a/scrapy/trunk/scrapy/utils/datatypes.py +++ b/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 - """ + """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)