@@ -25,6 +25,10 @@ class Full(Exception):
2525pass
2626272728+classShutDown(Exception):
29+'''Raised when put/get with shut-down queue.'''
30+31+2832classQueue:
2933'''Create a queue object with a given maximum size.
3034@@ -54,6 +58,9 @@ def __init__(self, maxsize=0):
5458self.all_tasks_done=threading.Condition(self.mutex)
5559self.unfinished_tasks=0
566061+# Queue shutdown state
62+self.is_shutdown=False
63+5764deftask_done(self):
5865'''Indicate that a formerly enqueued task is complete.
5966@@ -67,6 +74,8 @@ def task_done(self):
67746875 Raises a ValueError if called more times than there were items
6976 placed in the queue.
77+78+ Raises ShutDown if the queue has been shut down immediately.
7079 '''
7180withself.all_tasks_done:
7281unfinished=self.unfinished_tasks-1
@@ -84,6 +93,8 @@ def join(self):
8493 to indicate the item was retrieved and all work on it is complete.
85948695 When the count of unfinished tasks drops to zero, join() unblocks.
96+97+ Raises ShutDown if the queue has been shut down immediately.
8798 '''
8899withself.all_tasks_done:
89100whileself.unfinished_tasks:
@@ -129,15 +140,21 @@ def put(self, item, block=True, timeout=None):
129140 Otherwise ('block' is false), put an item on the queue if a free slot
130141 is immediately available, else raise the Full exception ('timeout'
131142 is ignored in that case).
143+144+ Raises ShutDown if the queue has been shut down.
132145 '''
133146withself.not_full:
147+ifself.is_shutdown:
148+raiseShutDown
134149ifself.maxsize>0:
135150ifnotblock:
136151ifself._qsize() >=self.maxsize:
137152raiseFull
138153eliftimeoutisNone:
139154whileself._qsize() >=self.maxsize:
140155self.not_full.wait()
156+ifself.is_shutdown:
157+raiseShutDown
141158eliftimeout<0:
142159raiseValueError("'timeout' must be a non-negative number")
143160else:
@@ -147,6 +164,8 @@ def put(self, item, block=True, timeout=None):
147164ifremaining<=0.0:
148165raiseFull
149166self.not_full.wait(remaining)
167+ifself.is_shutdown:
168+raiseShutDown
150169self._put(item)
151170self.unfinished_tasks+=1
152171self.not_empty.notify()
@@ -161,14 +180,21 @@ def get(self, block=True, timeout=None):
161180 Otherwise ('block' is false), return an item if one is immediately
162181 available, else raise the Empty exception ('timeout' is ignored
163182 in that case).
183+184+ Raises ShutDown if the queue has been shut down and is empty,
185+ or if the queue has been shut down immediately.
164186 '''
165187withself.not_empty:
188+ifself.is_shutdownandnotself._qsize():
189+raiseShutDown
166190ifnotblock:
167191ifnotself._qsize():
168192raiseEmpty
169193eliftimeoutisNone:
170194whilenotself._qsize():
171195self.not_empty.wait()
196+ifself.is_shutdownandnotself._qsize():
197+raiseShutDown
172198eliftimeout<0:
173199raiseValueError("'timeout' must be a non-negative number")
174200else:
@@ -178,6 +204,8 @@ def get(self, block=True, timeout=None):
178204ifremaining<=0.0:
179205raiseEmpty
180206self.not_empty.wait(remaining)
207+ifself.is_shutdownandnotself._qsize():
208+raiseShutDown
181209item=self._get()
182210self.not_full.notify()
183211returnitem
@@ -198,6 +226,28 @@ def get_nowait(self):
198226 '''
199227returnself.get(block=False)
200228229+defshutdown(self, immediate=False):
230+'''Shut-down the queue, making queue gets and puts raise.
231+232+ By default, gets will only raise once the queue is empty. Set
233+ 'immediate' to True to make gets raise immediately instead.
234+235+ All blocked callers of put() will be unblocked, and also get()
236+ and join() if 'immediate'. The ShutDown exception is raised.
237+ '''
238+withself.mutex:
239+self.is_shutdown=True
240+ifimmediate:
241+n_items=self._qsize()
242+whileself._qsize():
243+self._get()
244+ifself.unfinished_tasks>0:
245+self.unfinished_tasks-=1
246+self.not_empty.notify_all()
247+# release all blocked threads in `join()`
248+self.all_tasks_done.notify_all()
249+self.not_full.notify_all()
250+201251# Override these methods to implement other queue organizations
202252# (e.g. stack or priority queue).
203253# These will only be called with appropriate locks held