Lib/queue.py

   1 """A multi-producer, multi-consumer queue."""
   2
   3 from time import time as _time
   4 try:
   5     import threading as _threading
   6 except ImportError:
   7     import dummy_threading as _threading
   8 from collections import deque
   9 import heapq
  10
  11 __all__ = ['Empty', 'Full', 'Queue', 'PriorityQueue', 'LifoQueue']
  12
  13 class Empty(Exception):
  14     "Exception raised by Queue.get(block=0)/get_nowait()."
  15     pass
  16
  17 class Full(Exception):
  18     "Exception raised by Queue.put(block=0)/put_nowait()."
  19     pass
  20
  21 class Queue:
  22     """Create a queue object with a given maximum size.
  23
  24     If maxsize is <= 0, the queue size is infinite.
  25     """
  26     def __init__(self, maxsize=0):
  27         self.maxsize = maxsize
  28         self._init(maxsize)
  29         # mutex must be held whenever the queue is mutating.  All methods
  30         # that acquire mutex must release it before returning.  mutex
  31         # is shared between the three conditions, so acquiring and
  32         # releasing the conditions also acquires and releases mutex.
  33         self.mutex = _threading.Lock()
  34         # Notify not_empty whenever an item is added to the queue; a
  35         # thread waiting to get is notified then.
  36         self.not_empty = _threading.Condition(self.mutex)
  37         # Notify not_full whenever an item is removed from the queue;
  38         # a thread waiting to put is notified then.
  39         self.not_full = _threading.Condition(self.mutex)
  40         # Notify all_tasks_done whenever the number of unfinished tasks
  41         # drops to zero; thread waiting to join() is notified to resume
  42         self.all_tasks_done = _threading.Condition(self.mutex)
  43         self.unfinished_tasks = 0
  44
  45     def task_done(self):
  46         """Indicate that a formerly enqueued task is complete.
  47
  48         Used by Queue consumer threads.  For each get() used to fetch a task,
  49         a subsequent call to task_done() tells the queue that the processing
  50         on the task is complete.
  51
  52         If a join() is currently blocking, it will resume when all items
  53         have been processed (meaning that a task_done() call was received
  54         for every item that had been put() into the queue).
  55
  56         Raises a ValueError if called more times than there were items
  57         placed in the queue.
  58         """
  59         self.all_tasks_done.acquire()
  60         try:
  61             unfinished = self.unfinished_tasks - 1
  62             if unfinished <= 0:
  63                 if unfinished < 0:
  64                     raise ValueError('task_done() called too many times')
  65                 self.all_tasks_done.notify_all()
  66             self.unfinished_tasks = unfinished
  67         finally:
  68             self.all_tasks_done.release()
  69
  70     def join(self):
  71         """Blocks until all items in the Queue have been gotten and processed.
  72
  73         The count of unfinished tasks goes up whenever an item is added to the
  74         queue. The count goes down whenever a consumer thread calls task_done()
  75         to indicate the item was retrieved and all work on it is complete.
  76
  77         When the count of unfinished tasks drops to zero, join() unblocks.
  78         """
  79         self.all_tasks_done.acquire()
  80         try:
  81             while self.unfinished_tasks:
  82                 self.all_tasks_done.wait()
  83         finally:
  84             self.all_tasks_done.release()
  85
  86     def qsize(self):
  87         """Return the approximate size of the queue (not reliable!)."""
  88         self.mutex.acquire()
  89         n = self._qsize()
  90         self.mutex.release()
  91         return n
  92
  93     def empty(self):
  94         """Return True if the queue is empty, False otherwise (not reliable!).
  95
  96         This method is likely to be removed at some point.  Use qsize() == 0
  97         as a direct substitute, but be aware that either approach risks a race
  98         condition where a queue can grow before the result of empty() or
  99         qsize() can be used.
 100
 101         To create code that needs to wait for all queued tasks to be
 102         completed, the preferred technique is to use the join() method.
 103
 104         """
 105         self.mutex.acquire()
 106         n = not self._qsize()
 107         self.mutex.release()
 108         return n
 109
 110     def full(self):
 111         """Return True if the queue is full, False otherwise (not reliable!).
 112
 113         This method is likely to be removed at some point.  Use qsize() == n
 114         as a direct substitute, but be aware that either approach risks a race
 115         condition where a queue can shrink before the result of full() or
 116         qsize() can be used.
 117
 118         """
 119         self.mutex.acquire()
 120         n = 0 < self.maxsize == self._qsize()
 121         self.mutex.release()
 122         return n
 123
 124     def put(self, item, block=True, timeout=None):
 125         """Put an item into the queue.
 126
 127         If optional args 'block' is true and 'timeout' is None (the default),
 128         block if necessary until a free slot is available. If 'timeout' is
 129         a positive number, it blocks at most 'timeout' seconds and raises
 130         the Full exception if no free slot was available within that time.
 131         Otherwise ('block' is false), put an item on the queue if a free slot
 132         is immediately available, else raise the Full exception ('timeout'
 133         is ignored in that case).
 134         """
 135         self.not_full.acquire()
 136         try:
 137             if self.maxsize > 0:
 138                 if not block:
 139                     if self._qsize() == self.maxsize:
 140                         raise Full
 141                 elif timeout is None:
 142                     while self._qsize() == self.maxsize:
 143                         self.not_full.wait()
 144                 elif timeout < 0:
 145                     raise ValueError("'timeout' must be a positive number")
 146                 else:
 147                     endtime = _time() + timeout
 148                     while self._qsize() == self.maxsize:
 149                         remaining = endtime - _time()
 150                         if remaining <= 0.0:
 151                             raise Full
 152                         self.not_full.wait(remaining)
 153             self._put(item)
 154             self.unfinished_tasks += 1
 155             self.not_empty.notify()
 156         finally:
 157             self.not_full.release()
 158
 159     def put_nowait(self, item):
 160         """Put an item into the queue without blocking.
 161
 162         Only enqueue the item if a free slot is immediately available.
 163         Otherwise raise the Full exception.
 164         """
 165         return self.put(item, False)
 166
 167     def get(self, block=True, timeout=None):
 168         """Remove and return an item from the queue.
 169
 170         If optional args 'block' is true and 'timeout' is None (the default),
 171         block if necessary until an item is available. If 'timeout' is
 172         a positive number, it blocks at most 'timeout' seconds and raises
 173         the Empty exception if no item was available within that time.
 174         Otherwise ('block' is false), return an item if one is immediately
 175         available, else raise the Empty exception ('timeout' is ignored
 176         in that case).
 177         """
 178         self.not_empty.acquire()
 179         try:
 180             if not block:
 181                 if not self._qsize():
 182                     raise Empty
 183             elif timeout is None:
 184                 while not self._qsize():
 185                     self.not_empty.wait()
 186             elif timeout < 0:
 187                 raise ValueError("'timeout' must be a positive number")
 188             else:
 189                 endtime = _time() + timeout
 190                 while not self._qsize():
 191                     remaining = endtime - _time()
 192                     if remaining <= 0.0:
 193                         raise Empty
 194                     self.not_empty.wait(remaining)
 195             item = self._get()
 196             self.not_full.notify()
 197             return item
 198         finally:
 199             self.not_empty.release()
 200
 201     def get_nowait(self):
 202         """Remove and return an item from the queue without blocking.
 203
 204         Only get an item if one is immediately available. Otherwise
 205         raise the Empty exception.
 206         """
 207         return self.get(False)
 208
 209     # Override these methods to implement other queue organizations
 210     # (e.g. stack or priority queue).
 211     # These will only be called with appropriate locks held
 212
 213     # Initialize the queue representation
 214     def _init(self, maxsize):
 215         self.queue = deque()
 216
 217     def _qsize(self, len=len):
 218         return len(self.queue)
 219
 220     # Put a new item in the queue
 221     def _put(self, item):
 222         self.queue.append(item)
 223
 224     # Get an item from the queue
 225     def _get(self):
 226         return self.queue.popleft()
 227
 228
 229 class PriorityQueue(Queue):
 230     '''Variant of Queue that retrieves open entries in priority order (lowest first).
 231
 232     Entries are typically tuples of the form:  (priority number, data).
 233     '''
 234
 235     def _init(self, maxsize):
 236         self.queue = []
 237
 238     def _qsize(self, len=len):
 239         return len(self.queue)
 240
 241     def _put(self, item, heappush=heapq.heappush):
 242         heappush(self.queue, item)
 243
 244     def _get(self, heappop=heapq.heappop):
 245         return heappop(self.queue)
 246
 247
 248 class LifoQueue(Queue):
 249     '''Variant of Queue that retrieves most recently added entries first.'''
 250
 251     def _init(self, maxsize):
 252         self.queue = []
 253
 254     def _qsize(self, len=len):
 255         return len(self.queue)
 256
 257     def _put(self, item):
 258         self.queue.append(item)
 259
 260     def _get(self):
 261         return self.queue.pop()