blob: f051f1c89f2a879ef3bb6656f44eba433d270826 [file] [log] [blame]
Guido van Rossum4acc25b2000-02-02 15:10:15 +00001"""A multi-producer, multi-consumer queue."""
Guido van Rossum9022fce1992-08-25 12:30:44 +00002
Tim Peters5af0e412004-07-12 00:45:14 +00003from time import time as _time
Benjamin Peterson0ed52452009-03-21 17:36:10 +00004try:
5 import threading as _threading
6except ImportError:
7 import dummy_threading as _threading
Raymond Hettinger756b3f32004-01-29 06:37:52 +00008from collections import deque
Raymond Hettinger35641462008-01-17 00:13:27 +00009import heapq
Martin v. Löwis77ac4292002-10-15 15:11:13 +000010
Raymond Hettinger35641462008-01-17 00:13:27 +000011__all__ = ['Empty', 'Full', 'Queue', 'PriorityQueue', 'LifoQueue']
Brett Cannonb42bb5a2003-07-01 05:34:27 +000012
Tim Petersb8c94172001-01-15 22:53:46 +000013class Empty(Exception):
14 "Exception raised by Queue.get(block=0)/get_nowait()."
15 pass
16
17class Full(Exception):
18 "Exception raised by Queue.put(block=0)/put_nowait()."
19 pass
Guido van Rossum9022fce1992-08-25 12:30:44 +000020
21class Queue:
Thomas Wouters0e3f5912006-08-11 14:57:12 +000022 """Create a queue object with a given maximum size.
Guido van Rossum9022fce1992-08-25 12:30:44 +000023
Thomas Wouters0e3f5912006-08-11 14:57:12 +000024 If maxsize is <= 0, the queue size is infinite.
25 """
26 def __init__(self, maxsize=0):
Raymond Hettingerda3caed2008-01-14 21:39:24 +000027 self.maxsize = maxsize
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000028 self._init(maxsize)
Tim Peters5af0e412004-07-12 00:45:14 +000029 # mutex must be held whenever the queue is mutating. All methods
30 # that acquire mutex must release it before returning. mutex
Thomas Wouters89f507f2006-12-13 04:49:30 +000031 # is shared between the three conditions, so acquiring and
Tim Peters5af0e412004-07-12 00:45:14 +000032 # releasing the conditions also acquires and releases mutex.
Benjamin Peterson0ed52452009-03-21 17:36:10 +000033 self.mutex = _threading.Lock()
Tim Peters5af0e412004-07-12 00:45:14 +000034 # Notify not_empty whenever an item is added to the queue; a
35 # thread waiting to get is notified then.
Benjamin Peterson0ed52452009-03-21 17:36:10 +000036 self.not_empty = _threading.Condition(self.mutex)
Tim Peters5af0e412004-07-12 00:45:14 +000037 # Notify not_full whenever an item is removed from the queue;
38 # a thread waiting to put is notified then.
Benjamin Peterson0ed52452009-03-21 17:36:10 +000039 self.not_full = _threading.Condition(self.mutex)
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000040 # Notify all_tasks_done whenever the number of unfinished tasks
41 # drops to zero; thread waiting to join() is notified to resume
Benjamin Peterson0ed52452009-03-21 17:36:10 +000042 self.all_tasks_done = _threading.Condition(self.mutex)
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000043 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')
Benjamin Peterson672b8032008-06-11 19:14:14 +000065 self.all_tasks_done.notify_all()
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000066 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()
Guido van Rossum9022fce1992-08-25 12:30:44 +000085
Barry Warsaw3d96d521997-11-20 19:56:38 +000086 def qsize(self):
Guido van Rossum9e1721f1999-02-08 18:34:01 +000087 """Return the approximate size of the queue (not reliable!)."""
Guido van Rossum7e6d18c1998-04-29 14:29:32 +000088 self.mutex.acquire()
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000089 n = self._qsize()
Guido van Rossum7e6d18c1998-04-29 14:29:32 +000090 self.mutex.release()
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000091 return n
Barry Warsaw3d96d521997-11-20 19:56:38 +000092
Alexandre Vassalottif260e442008-05-11 19:59:59 +000093 def empty(self):
Raymond Hettinger611eaf02009-03-06 23:55:28 +000094 """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 """
Alexandre Vassalottif260e442008-05-11 19:59:59 +0000105 self.mutex.acquire()
106 n = not self._qsize()
107 self.mutex.release()
108 return n
109
110 def full(self):
Raymond Hettinger611eaf02009-03-06 23:55:28 +0000111 """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 """
Alexandre Vassalottif260e442008-05-11 19:59:59 +0000119 self.mutex.acquire()
120 n = 0 < self.maxsize == self._qsize()
121 self.mutex.release()
122 return n
123
Martin v. Löwis77ac4292002-10-15 15:11:13 +0000124 def put(self, item, block=True, timeout=None):
Guido van Rossumc09e6b11998-04-09 14:25:32 +0000125 """Put an item into the queue.
126
Martin v. Löwis77ac4292002-10-15 15:11:13 +0000127 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).
Guido van Rossum9e1721f1999-02-08 18:34:01 +0000134 """
Tim Peters5af0e412004-07-12 00:45:14 +0000135 self.not_full.acquire()
Mark Hammond3b959db2002-04-19 00:11:32 +0000136 try:
Raymond Hettingerae138cb2008-01-15 20:42:00 +0000137 if self.maxsize > 0:
138 if not block:
139 if self._qsize() == self.maxsize:
140 raise Full
141 elif timeout is None:
Raymond Hettingerda3caed2008-01-14 21:39:24 +0000142 while self._qsize() == self.maxsize:
143 self.not_full.wait()
Raymond Hettingerae138cb2008-01-15 20:42:00 +0000144 elif timeout < 0:
Tim Peters5af0e412004-07-12 00:45:14 +0000145 raise ValueError("'timeout' must be a positive number")
Raymond Hettingerae138cb2008-01-15 20:42:00 +0000146 else:
147 endtime = _time() + timeout
Raymond Hettingerda3caed2008-01-14 21:39:24 +0000148 while self._qsize() == self.maxsize:
149 remaining = endtime - _time()
150 if remaining <= 0.0:
151 raise Full
152 self.not_full.wait(remaining)
Mark Hammond3b959db2002-04-19 00:11:32 +0000153 self._put(item)
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000154 self.unfinished_tasks += 1
Tim Peters5af0e412004-07-12 00:45:14 +0000155 self.not_empty.notify()
Mark Hammond3b959db2002-04-19 00:11:32 +0000156 finally:
Tim Peters5af0e412004-07-12 00:45:14 +0000157 self.not_full.release()
Barry Warsaw3d96d521997-11-20 19:56:38 +0000158
Guido van Rossum9e1721f1999-02-08 18:34:01 +0000159 def put_nowait(self, item):
160 """Put an item into the queue without blocking.
Guido van Rossumc09e6b11998-04-09 14:25:32 +0000161
Guido van Rossum9e1721f1999-02-08 18:34:01 +0000162 Only enqueue the item if a free slot is immediately available.
163 Otherwise raise the Full exception.
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000164 """
Tim Peters71ed2202004-07-12 01:20:32 +0000165 return self.put(item, False)
Barry Warsaw3d96d521997-11-20 19:56:38 +0000166
Martin v. Löwis77ac4292002-10-15 15:11:13 +0000167 def get(self, block=True, timeout=None):
Guido van Rossum9e1721f1999-02-08 18:34:01 +0000168 """Remove and return an item from the queue.
Guido van Rossumc09e6b11998-04-09 14:25:32 +0000169
Martin v. Löwis77ac4292002-10-15 15:11:13 +0000170 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).
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000177 """
Tim Peters5af0e412004-07-12 00:45:14 +0000178 self.not_empty.acquire()
Mark Hammond3b959db2002-04-19 00:11:32 +0000179 try:
Tim Peters71ed2202004-07-12 01:20:32 +0000180 if not block:
Raymond Hettingerda3caed2008-01-14 21:39:24 +0000181 if not self._qsize():
Tim Peters71ed2202004-07-12 01:20:32 +0000182 raise Empty
183 elif timeout is None:
Raymond Hettingerda3caed2008-01-14 21:39:24 +0000184 while not self._qsize():
Tim Peters5af0e412004-07-12 00:45:14 +0000185 self.not_empty.wait()
Raymond Hettingerae138cb2008-01-15 20:42:00 +0000186 elif timeout < 0:
187 raise ValueError("'timeout' must be a positive number")
Tim Peters5af0e412004-07-12 00:45:14 +0000188 else:
Tim Peters5af0e412004-07-12 00:45:14 +0000189 endtime = _time() + timeout
Raymond Hettingerda3caed2008-01-14 21:39:24 +0000190 while not self._qsize():
Tim Peters5af0e412004-07-12 00:45:14 +0000191 remaining = endtime - _time()
Tim Peters71ed2202004-07-12 01:20:32 +0000192 if remaining <= 0.0:
Tim Peters5af0e412004-07-12 00:45:14 +0000193 raise Empty
194 self.not_empty.wait(remaining)
Mark Hammond3b959db2002-04-19 00:11:32 +0000195 item = self._get()
Tim Peters5af0e412004-07-12 00:45:14 +0000196 self.not_full.notify()
197 return item
Mark Hammond3b959db2002-04-19 00:11:32 +0000198 finally:
Tim Peters5af0e412004-07-12 00:45:14 +0000199 self.not_empty.release()
Guido van Rossum9022fce1992-08-25 12:30:44 +0000200
Guido van Rossum9e1721f1999-02-08 18:34:01 +0000201 def get_nowait(self):
202 """Remove and return an item from the queue without blocking.
Guido van Rossum9022fce1992-08-25 12:30:44 +0000203
Martin v. Löwis77ac4292002-10-15 15:11:13 +0000204 Only get an item if one is immediately available. Otherwise
Guido van Rossum9e1721f1999-02-08 18:34:01 +0000205 raise the Empty exception.
206 """
Tim Peters71ed2202004-07-12 01:20:32 +0000207 return self.get(False)
Guido van Rossum9022fce1992-08-25 12:30:44 +0000208
Barry Warsaw3d96d521997-11-20 19:56:38 +0000209 # 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
Guido van Rossum9022fce1992-08-25 12:30:44 +0000212
Barry Warsaw3d96d521997-11-20 19:56:38 +0000213 # Initialize the queue representation
214 def _init(self, maxsize):
Raymond Hettinger756b3f32004-01-29 06:37:52 +0000215 self.queue = deque()
Guido van Rossum9022fce1992-08-25 12:30:44 +0000216
Raymond Hettinger35641462008-01-17 00:13:27 +0000217 def _qsize(self, len=len):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000218 return len(self.queue)
Guido van Rossum9022fce1992-08-25 12:30:44 +0000219
Barry Warsaw3d96d521997-11-20 19:56:38 +0000220 # Put a new item in the queue
221 def _put(self, item):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000222 self.queue.append(item)
Guido van Rossum9022fce1992-08-25 12:30:44 +0000223
Barry Warsaw3d96d521997-11-20 19:56:38 +0000224 # Get an item from the queue
225 def _get(self):
Raymond Hettinger756b3f32004-01-29 06:37:52 +0000226 return self.queue.popleft()
Raymond Hettinger35641462008-01-17 00:13:27 +0000227
228
229class 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
248class 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()