@@ -25,6 +25,10 @@ class Full(Exception):
25
25
pass
26
26
27
27
28
+ class ShutDown (Exception ):
29
+ '''Raised when put/get with shut-down queue.'''
30
+
31
+
28
32
class Queue :
29
33
'''Create a queue object with a given maximum size.
30
34
@@ -54,6 +58,9 @@ def __init__(self, maxsize=0):
54
58
self .all_tasks_done = threading .Condition (self .mutex )
55
59
self .unfinished_tasks = 0
56
60
61
+ # Queue shutdown state
62
+ self .is_shutdown = False
63
+
57
64
def task_done (self ):
58
65
'''Indicate that a formerly enqueued task is complete.
59
66
@@ -67,6 +74,8 @@ def task_done(self):
67
74
68
75
Raises a ValueError if called more times than there were items
69
76
placed in the queue.
77
+
78
+ Raises ShutDown if the queue has been shut down immediately.
70
79
'''
71
80
with self .all_tasks_done :
72
81
unfinished = self .unfinished_tasks - 1
@@ -84,6 +93,8 @@ def join(self):
84
93
to indicate the item was retrieved and all work on it is complete.
85
94
86
95
When the count of unfinished tasks drops to zero, join() unblocks.
96
+
97
+ Raises ShutDown if the queue has been shut down immediately.
87
98
'''
88
99
with self .all_tasks_done :
89
100
while self .unfinished_tasks :
@@ -129,15 +140,21 @@ def put(self, item, block=True, timeout=None):
129
140
Otherwise ('block' is false), put an item on the queue if a free slot
130
141
is immediately available, else raise the Full exception ('timeout'
131
142
is ignored in that case).
143
+
144
+ Raises ShutDown if the queue has been shut down.
132
145
'''
133
146
with self .not_full :
147
+ if self .is_shutdown :
148
+ raise ShutDown
134
149
if self .maxsize > 0 :
135
150
if not block :
136
151
if self ._qsize () >= self .maxsize :
137
152
raise Full
138
153
elif timeout is None :
139
154
while self ._qsize () >= self .maxsize :
140
155
self .not_full .wait ()
156
+ if self .is_shutdown :
157
+ raise ShutDown
141
158
elif timeout < 0 :
142
159
raise ValueError ("'timeout' must be a non-negative number" )
143
160
else :
@@ -147,6 +164,8 @@ def put(self, item, block=True, timeout=None):
147
164
if remaining <= 0.0 :
148
165
raise Full
149
166
self .not_full .wait (remaining )
167
+ if self .is_shutdown :
168
+ raise ShutDown
150
169
self ._put (item )
151
170
self .unfinished_tasks += 1
152
171
self .not_empty .notify ()
@@ -161,14 +180,21 @@ def get(self, block=True, timeout=None):
161
180
Otherwise ('block' is false), return an item if one is immediately
162
181
available, else raise the Empty exception ('timeout' is ignored
163
182
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.
164
186
'''
165
187
with self .not_empty :
188
+ if self .is_shutdown and not self ._qsize ():
189
+ raise ShutDown
166
190
if not block :
167
191
if not self ._qsize ():
168
192
raise Empty
169
193
elif timeout is None :
170
194
while not self ._qsize ():
171
195
self .not_empty .wait ()
196
+ if self .is_shutdown and not self ._qsize ():
197
+ raise ShutDown
172
198
elif timeout < 0 :
173
199
raise ValueError ("'timeout' must be a non-negative number" )
174
200
else :
@@ -178,6 +204,8 @@ def get(self, block=True, timeout=None):
178
204
if remaining <= 0.0 :
179
205
raise Empty
180
206
self .not_empty .wait (remaining )
207
+ if self .is_shutdown and not self ._qsize ():
208
+ raise ShutDown
181
209
item = self ._get ()
182
210
self .not_full .notify ()
183
211
return item
@@ -198,6 +226,28 @@ def get_nowait(self):
198
226
'''
199
227
return self .get (block = False )
200
228
229
+ def shutdown (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
+ with self .mutex :
239
+ self .is_shutdown = True
240
+ if immediate :
241
+ n_items = self ._qsize ()
242
+ while self ._qsize ():
243
+ self ._get ()
244
+ if self .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
+
201
251
# Override these methods to implement other queue organizations
202
252
# (e.g. stack or priority queue).
203
253
# These will only be called with appropriate locks held
0 commit comments