[81] | 1 | from heapq import heappush, heappop
|
---|
| 2 | from queue import Queue
|
---|
| 3 | from typing import TypeVar, Generic, Callable, List, Set
|
---|
| 4 |
|
---|
| 5 |
|
---|
| 6 | T=TypeVar("T")
|
---|
| 7 | class Wrap(Generic[T]):
|
---|
| 8 | '''
|
---|
| 9 | Wrapper that has __lt__.
|
---|
| 10 | We wrap objects that possibly don't implement __lt__
|
---|
| 11 | in the queue in this to avoid
|
---|
| 12 | heapq crashing on a call to __lt__
|
---|
| 13 | '''
|
---|
| 14 | def __init__(self, obj:T, lessthanfunction):
|
---|
| 15 | '''
|
---|
| 16 | @param obj the object to be wrapped.
|
---|
| 17 | @param lessthanfunction a function that
|
---|
| 18 | when called with another object,
|
---|
| 19 | returns true iff other object is < this.
|
---|
| 20 | '''
|
---|
| 21 | self._obj=obj
|
---|
| 22 | self._lessthanfunction = lessthanfunction
|
---|
| 23 | def get(self) -> T:
|
---|
| 24 | return self._obj
|
---|
| 25 | def __repr__(self):
|
---|
| 26 | return str(self._obj)
|
---|
| 27 | def __lt__(self, other):
|
---|
| 28 | return self._lessthanfunction(self.get(), other.get())
|
---|
| 29 |
|
---|
| 30 | class PriorityQueue(Queue, Generic[T]):
|
---|
| 31 | '''Variant of Queue that retrieves open entries in priority order (lowest first).
|
---|
| 32 | This differs from queue.PriorityQueue as it allows custom comparator function.
|
---|
| 33 | Another difference is that entries are
|
---|
| 34 | '''
|
---|
| 35 |
|
---|
| 36 | def __init__(self, lessthanfunction:Callable[[T,T], bool]):
|
---|
| 37 | '''
|
---|
| 38 | @param lessthanfunction a function with arguments (this, other)
|
---|
| 39 | that is used to sort the queue.
|
---|
| 40 | If queue should have the "smallest" item at head,
|
---|
| 41 | this function should return true iff other object is < this.
|
---|
| 42 | '''
|
---|
| 43 | super().__init__()
|
---|
| 44 | self._queue:List[Wrap[T]] = []
|
---|
| 45 | self._lessthanfunction=lessthanfunction
|
---|
| 46 |
|
---|
| 47 | def _qsize(self)->int:
|
---|
| 48 | return len(self._queue)
|
---|
| 49 |
|
---|
| 50 | def _put(self, item:T):
|
---|
| 51 | heappush(self._queue, Wrap(item, self._lessthanfunction))
|
---|
| 52 |
|
---|
| 53 | def _get(self)->T:
|
---|
| 54 | return heappop(self._queue).get()
|
---|
| 55 |
|
---|
| 56 | def __repr__(self)->str:
|
---|
| 57 | return str(self._queue)
|
---|
| 58 |
|
---|
| 59 | def __iter__(self):
|
---|
| 60 | return iter([w.get() for w in self._queue])
|
---|
| 61 |
|
---|
| 62 | |
---|