class PriorityQueue(object): """A priority queue implemented using multiple internal queues (typically, FIFO queues). The internal queue must implement the following methods: * push(obj) * pop() * close() * __len__() The constructor receives a qfactory argument, which is a callable used to instantiate a new (internal) queue when a new priority is allocated. The qfactory function is called with the priority number as first and only argument. Only integer priorities should be used. Lower numbers are higher priorities. startprios is a sequence of priorities to start with. If the queue was previously closed leaving some priority buckets non-empty, those priorities should be passed in startprios. """ def __init__(self, qfactory, startprios=()): self.queues = {} self.qfactory = qfactory for p in startprios: self.queues[p] = self.qfactory(p) self.curprio = min(startprios) if startprios else None def push(self, obj, priority=0): if priority not in self.queues: self.queues[priority] = self.qfactory(priority) q = self.queues[priority] q.push(obj) # this may fail (eg. serialization error) if self.curprio is None or priority < self.curprio: self.curprio = priority def pop(self): if self.curprio is None: return q = self.queues[self.curprio] m = q.pop() if len(q) == 0: del self.queues[self.curprio] q.close() prios = [p for p, q in self.queues.items() if len(q) > 0] self.curprio = min(prios) if prios else None return m def close(self): active = [] for p, q in self.queues.items(): if len(q): active.append(p) q.close() return active def __len__(self): return sum(len(x) for x in self.queues.values()) if self.queues else 0