blob: 4d89d97c4534a90042514e73629c65012480b722 [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
Raymond Hettinger756b3f32004-01-29 06:37:52 +00004from collections import deque
Raymond Hettinger35641462008-01-17 00:13:27 +00005import heapq
Martin v. Löwis77ac4292002-10-15 15:11:13 +00006
Raymond Hettinger35641462008-01-17 00:13:27 +00007__all__ = ['Empty', 'Full', 'Queue', 'PriorityQueue', 'LifoQueue']
Brett Cannonb42bb5a2003-07-01 05:34:27 +00008
Tim Petersb8c94172001-01-15 22:53:46 +00009class Empty(Exception):
10 "Exception raised by Queue.get(block=0)/get_nowait()."
11 pass
12
13class Full(Exception):
14 "Exception raised by Queue.put(block=0)/put_nowait()."
15 pass
Guido van Rossum9022fce1992-08-25 12:30:44 +000016
17class Queue:
Thomas Wouters0e3f5912006-08-11 14:57:12 +000018 """Create a queue object with a given maximum size.
Guido van Rossum9022fce1992-08-25 12:30:44 +000019
Thomas Wouters0e3f5912006-08-11 14:57:12 +000020 If maxsize is <= 0, the queue size is infinite.
21 """
22 def __init__(self, maxsize=0):
Guido van Rossuma0934242002-12-30 22:36:09 +000023 try:
Tim Peters5af0e412004-07-12 00:45:14 +000024 import threading
Guido van Rossuma0934242002-12-30 22:36:09 +000025 except ImportError:
Tim Peters5af0e412004-07-12 00:45:14 +000026 import dummy_threading as threading
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.
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)
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
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.notifyAll()
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()
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
Martin v. Löwis77ac4292002-10-15 15:11:13 +000093 def put(self, item, block=True, timeout=None):
Guido van Rossumc09e6b11998-04-09 14:25:32 +000094 """Put an item into the queue.
95
Martin v. Löwis77ac4292002-10-15 15:11:13 +000096 If optional args 'block' is true and 'timeout' is None (the default),
97 block if necessary until a free slot is available. If 'timeout' is
98 a positive number, it blocks at most 'timeout' seconds and raises
99 the Full exception if no free slot was available within that time.
100 Otherwise ('block' is false), put an item on the queue if a free slot
101 is immediately available, else raise the Full exception ('timeout'
102 is ignored in that case).
Guido van Rossum9e1721f1999-02-08 18:34:01 +0000103 """
Tim Peters5af0e412004-07-12 00:45:14 +0000104 self.not_full.acquire()
Mark Hammond3b959db2002-04-19 00:11:32 +0000105 try:
Raymond Hettingerae138cb2008-01-15 20:42:00 +0000106 if self.maxsize > 0:
107 if not block:
108 if self._qsize() == self.maxsize:
109 raise Full
110 elif timeout is None:
Raymond Hettingerda3caed2008-01-14 21:39:24 +0000111 while self._qsize() == self.maxsize:
112 self.not_full.wait()
Raymond Hettingerae138cb2008-01-15 20:42:00 +0000113 elif timeout < 0:
Tim Peters5af0e412004-07-12 00:45:14 +0000114 raise ValueError("'timeout' must be a positive number")
Raymond Hettingerae138cb2008-01-15 20:42:00 +0000115 else:
116 endtime = _time() + timeout
Raymond Hettingerda3caed2008-01-14 21:39:24 +0000117 while self._qsize() == self.maxsize:
118 remaining = endtime - _time()
119 if remaining <= 0.0:
120 raise Full
121 self.not_full.wait(remaining)
Mark Hammond3b959db2002-04-19 00:11:32 +0000122 self._put(item)
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000123 self.unfinished_tasks += 1
Tim Peters5af0e412004-07-12 00:45:14 +0000124 self.not_empty.notify()
Mark Hammond3b959db2002-04-19 00:11:32 +0000125 finally:
Tim Peters5af0e412004-07-12 00:45:14 +0000126 self.not_full.release()
Barry Warsaw3d96d521997-11-20 19:56:38 +0000127
Guido van Rossum9e1721f1999-02-08 18:34:01 +0000128 def put_nowait(self, item):
129 """Put an item into the queue without blocking.
Guido van Rossumc09e6b11998-04-09 14:25:32 +0000130
Guido van Rossum9e1721f1999-02-08 18:34:01 +0000131 Only enqueue the item if a free slot is immediately available.
132 Otherwise raise the Full exception.
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000133 """
Tim Peters71ed2202004-07-12 01:20:32 +0000134 return self.put(item, False)
Barry Warsaw3d96d521997-11-20 19:56:38 +0000135
Martin v. Löwis77ac4292002-10-15 15:11:13 +0000136 def get(self, block=True, timeout=None):
Guido van Rossum9e1721f1999-02-08 18:34:01 +0000137 """Remove and return an item from the queue.
Guido van Rossumc09e6b11998-04-09 14:25:32 +0000138
Martin v. Löwis77ac4292002-10-15 15:11:13 +0000139 If optional args 'block' is true and 'timeout' is None (the default),
140 block if necessary until an item is available. If 'timeout' is
141 a positive number, it blocks at most 'timeout' seconds and raises
142 the Empty exception if no item was available within that time.
143 Otherwise ('block' is false), return an item if one is immediately
144 available, else raise the Empty exception ('timeout' is ignored
145 in that case).
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000146 """
Tim Peters5af0e412004-07-12 00:45:14 +0000147 self.not_empty.acquire()
Mark Hammond3b959db2002-04-19 00:11:32 +0000148 try:
Tim Peters71ed2202004-07-12 01:20:32 +0000149 if not block:
Raymond Hettingerda3caed2008-01-14 21:39:24 +0000150 if not self._qsize():
Tim Peters71ed2202004-07-12 01:20:32 +0000151 raise Empty
152 elif timeout is None:
Raymond Hettingerda3caed2008-01-14 21:39:24 +0000153 while not self._qsize():
Tim Peters5af0e412004-07-12 00:45:14 +0000154 self.not_empty.wait()
Raymond Hettingerae138cb2008-01-15 20:42:00 +0000155 elif timeout < 0:
156 raise ValueError("'timeout' must be a positive number")
Tim Peters5af0e412004-07-12 00:45:14 +0000157 else:
Tim Peters5af0e412004-07-12 00:45:14 +0000158 endtime = _time() + timeout
Raymond Hettingerda3caed2008-01-14 21:39:24 +0000159 while not self._qsize():
Tim Peters5af0e412004-07-12 00:45:14 +0000160 remaining = endtime - _time()
Tim Peters71ed2202004-07-12 01:20:32 +0000161 if remaining <= 0.0:
Tim Peters5af0e412004-07-12 00:45:14 +0000162 raise Empty
163 self.not_empty.wait(remaining)
Mark Hammond3b959db2002-04-19 00:11:32 +0000164 item = self._get()
Tim Peters5af0e412004-07-12 00:45:14 +0000165 self.not_full.notify()
166 return item
Mark Hammond3b959db2002-04-19 00:11:32 +0000167 finally:
Tim Peters5af0e412004-07-12 00:45:14 +0000168 self.not_empty.release()
Guido van Rossum9022fce1992-08-25 12:30:44 +0000169
Guido van Rossum9e1721f1999-02-08 18:34:01 +0000170 def get_nowait(self):
171 """Remove and return an item from the queue without blocking.
Guido van Rossum9022fce1992-08-25 12:30:44 +0000172
Martin v. Löwis77ac4292002-10-15 15:11:13 +0000173 Only get an item if one is immediately available. Otherwise
Guido van Rossum9e1721f1999-02-08 18:34:01 +0000174 raise the Empty exception.
175 """
Tim Peters71ed2202004-07-12 01:20:32 +0000176 return self.get(False)
Guido van Rossum9022fce1992-08-25 12:30:44 +0000177
Barry Warsaw3d96d521997-11-20 19:56:38 +0000178 # Override these methods to implement other queue organizations
179 # (e.g. stack or priority queue).
180 # These will only be called with appropriate locks held
Guido van Rossum9022fce1992-08-25 12:30:44 +0000181
Barry Warsaw3d96d521997-11-20 19:56:38 +0000182 # Initialize the queue representation
183 def _init(self, maxsize):
Raymond Hettinger756b3f32004-01-29 06:37:52 +0000184 self.queue = deque()
Guido van Rossum9022fce1992-08-25 12:30:44 +0000185
Raymond Hettinger35641462008-01-17 00:13:27 +0000186 def _qsize(self, len=len):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000187 return len(self.queue)
Guido van Rossum9022fce1992-08-25 12:30:44 +0000188
Barry Warsaw3d96d521997-11-20 19:56:38 +0000189 # Put a new item in the queue
190 def _put(self, item):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000191 self.queue.append(item)
Guido van Rossum9022fce1992-08-25 12:30:44 +0000192
Barry Warsaw3d96d521997-11-20 19:56:38 +0000193 # Get an item from the queue
194 def _get(self):
Raymond Hettinger756b3f32004-01-29 06:37:52 +0000195 return self.queue.popleft()
Raymond Hettinger35641462008-01-17 00:13:27 +0000196
197
198class PriorityQueue(Queue):
199 '''Variant of Queue that retrieves open entries in priority order (lowest first).
200
201 Entries are typically tuples of the form: (priority number, data).
202 '''
203
204 def _init(self, maxsize):
205 self.queue = []
206
207 def _qsize(self, len=len):
208 return len(self.queue)
209
210 def _put(self, item, heappush=heapq.heappush):
211 heappush(self.queue, item)
212
213 def _get(self, heappop=heapq.heappop):
214 return heappop(self.queue)
215
216
217class LifoQueue(Queue):
218 '''Variant of Queue that retrieves most recently added entries first.'''
219
220 def _init(self, maxsize):
221 self.queue = []
222
223 def _qsize(self, len=len):
224 return len(self.queue)
225
226 def _put(self, item):
227 self.queue.append(item)
228
229 def _get(self):
230 return self.queue.pop()