new module containing a class to handle prioritized tasks queue

2 files changed, 128 insertions, 0 deletions

diff --git a/tasksqueue.py b/tasksqueue.py
new file mode 100644
index 0000000..75ade01
--- /dev/null
+++ b/tasksqueue.py
@@ -0,0 +1,74 @@
+"""Prioritized tasks queue
+
+:organization: Logilab
+:copyright: 2008 LOGILAB S.A. (Paris, FRANCE), all rights reserved.
+:contact: http://www.logilab.fr/ -- mailto:contact@logilab.fr
+"""
+__docformat__ = "restructuredtext en"
+
+from bisect import insort_left
+from Queue import Queue
+
+LOW = 0
+MEDIUM = 10
+HIGH = 100
+
+
+class PrioritizedTasksQueue(Queue):
+
+    def _init(self, maxsize):
+        """Initialize the queue representation"""
+        self.maxsize = maxsize
+        # ordered list of task, from the lowest to the highest priority
+        self.queue = []
+
+    def _put(self, item):
+        """Put a new item in the queue"""
+        for i, task in enumerate(self.queue):
+            # equivalent task
+            if task == item:
+                # if new task has a higher priority, remove the one already
+                # queued so the new priority will be considered
+                if task < item:
+                    item.merge(task)
+                    del self.queue[i]
+                    break
+                # else keep it so current order is kept
+                task.merge(item)
+                return
+        insort_left(self.queue, item)
+        
+    def _get(self):
+        """Get an item from the queue"""
+        return self.queue.pop()
+
+    def __iter__(self):
+        return iter(self.queue)
+
+    def remove(self, tid):
+        """remove a specific task from the queue"""
+        # XXX acquire lock
+        for i, task in enumerate(self):
+            if task.id == tid:
+                self.queue.pop(i)
+                return
+        raise ValueError('not task of id %s in queue' % tid)
+
+class Task(object):
+    def __init__(self, tid, priority=LOW):
+        # task id
+        self.id = tid
+        # task priority
+        self.priority = priority
+
+    def __repr__(self):
+        return '<Task %s @%#x>' % (self.id, id(self))
+    
+    def __cmp__(self, other):
+        return cmp(self.priority, other.priority)
+    
+    def __eq__(self, other):
+        return self.id == other.id
+
+    def merge(self, other):
+        pass
diff --git a/test/unittest_taskqueue.py b/test/unittest_taskqueue.py
new file mode 100644
index 0000000..a723d1c
--- /dev/null
+++ b/test/unittest_taskqueue.py
@@ -0,0 +1,54 @@
+from logilab.common.testlib import TestCase, unittest_main
+
+from logilab.common.tasksqueue import *
+
+class TaskTC(TestCase):
+    
+    def test_eq(self):
+        self.failIf(Task('t1') == Task('t2'))
+        self.failUnless(Task('t1') == Task('t1'))
+        
+    def test_cmp(self):
+        self.failUnless(Task('t1', LOW) < Task('t2', MEDIUM))
+        self.failIf(Task('t1', LOW) > Task('t2', MEDIUM))
+        self.failUnless(Task('t1', HIGH) > Task('t2', MEDIUM))
+        self.failIf(Task('t1', HIGH) < Task('t2', MEDIUM))
+
+        
+class PrioritizedTasksQueueTC(TestCase):
+    
+    def test_priority(self):
+        queue = PrioritizedTasksQueue()
+        queue.put(Task('t1'))
+        queue.put(Task('t2', MEDIUM))
+        queue.put(Task('t3', HIGH))
+        queue.put(Task('t4', LOW))
+        self.assertEquals(queue.get().id, 't3')
+        self.assertEquals(queue.get().id, 't2')
+        self.assertEquals(queue.get().id, 't1')
+        self.assertEquals(queue.get().id, 't4')
+        
+    def test_remove_equivalent(self):
+        queue = PrioritizedTasksQueue()
+        queue.put(Task('t1'))
+        queue.put(Task('t2', MEDIUM))
+        queue.put(Task('t1', HIGH))
+        queue.put(Task('t3', MEDIUM))
+        queue.put(Task('t2', MEDIUM))
+        self.assertEquals(queue.qsize(), 3)
+        self.assertEquals(queue.get().id, 't1')
+        self.assertEquals(queue.get().id, 't2')
+        self.assertEquals(queue.get().id, 't3')
+        self.assertEquals(queue.qsize(), 0)
+        
+    def test_remove(self):
+        queue = PrioritizedTasksQueue()
+        queue.put(Task('t1'))
+        queue.put(Task('t2'))
+        queue.put(Task('t3'))
+        queue.remove('t2')
+        self.assertEquals([t.id for t in queue], ['t3', 't1'])
+        self.assertRaises(ValueError, queue.remove, 't4')
+        
+if __name__ == '__main__':
+    unittest_main()