| 2 | from Queue import Queue |
| 3 | |
| 4 | class TaskQueue(Queue): |
| 5 | |
| 6 | def __init__(self): |
| 7 | Queue.__init__(self) |
| 8 | self.all_tasks_done = threading.Condition(self.mutex) |
| 9 | self.unfinished_tasks = 0 |
| 10 | |
| 11 | def _put(self, item): |
| 12 | Queue._put(self, item) |
| 13 | self.unfinished_tasks += 1 |
| 14 | |
| 15 | def task_done(self): |
| 16 | """Indicate that a formerly enqueued task is complete. |
| 17 | |
| 18 | Used by Queue consumer threads. For each get() used to fetch a task, |
| 19 | a subsequent call to task_done() tells the queue that the processing |
| 20 | on the task is complete. |
| 21 | |
| 22 | If a join() is currently blocking, it will resume when all items |
| 23 | have been processed (meaning that a task_done() call was received |
| 24 | for every item that had been put() into the queue). |
| 25 | |
| 26 | Raises a ValueError if called more times than there were items |
| 27 | placed in the queue. |
| 28 | """ |
| 29 | self.all_tasks_done.acquire() |
| 30 | try: |
| 31 | unfinished = self.unfinished_tasks - 1 |
| 32 | if unfinished <= 0: |
| 33 | if unfinished < 0: |
| 34 | raise ValueError('task_done() called too many times') |
| 35 | self.all_tasks_done.notifyAll() |
| 36 | self.unfinished_tasks = unfinished |
| 37 | finally: |
| 38 | self.all_tasks_done.release() |
| 39 | |
| 40 | def join(self): |
| 41 | """Blocks until all items in the Queue have been gotten and processed. |
| 42 | |
| 43 | The count of unfinished tasks goes up whenever an item is added to the |
| 44 | queue. The count goes down whenever a consumer thread calls task_done() |
| 45 | to indicate the item was retrieved and all work on it is complete. |
| 46 | |
| 47 | When the count of unfinished tasks drops to zero, join() unblocks. |
| 48 | """ |
| 49 | self.all_tasks_done.acquire() |
| 50 | try: |
| 51 | while self.unfinished_tasks: |
| 52 | self.all_tasks_done.wait() |
| 53 | finally: |
| 54 | self.all_tasks_done.release() |
| 55 | |
| 56 | |
| 57 | #### Example code #################### |