8
8
)
9
9
10
10
import collections
11
+ import enum
11
12
import heapq
12
13
from types import GenericAlias
13
14
@@ -30,9 +31,10 @@ class QueueShutDown(Exception):
30
31
pass
31
32
32
33
33
- _queue_alive = "alive"
34
- _queue_shutdown = "shutdown"
35
- _queue_shutdown_immediate = "shutdown-immediate"
34
+ class _QueueState (enum .Enum ):
35
+ ALIVE = "alive"
36
+ SHUTDOWN = "shutdown"
37
+ SHUTDOWN_IMMEDIATE = "shutdown-immediate"
36
38
37
39
38
40
class Queue (mixins ._LoopBoundMixin ):
@@ -58,7 +60,7 @@ def __init__(self, maxsize=0):
58
60
self ._finished = locks .Event ()
59
61
self ._finished .set ()
60
62
self ._init (maxsize )
61
- self .shutdown_state = _queue_alive
63
+ self ._shutdown_state = _QueueState . ALIVE
62
64
63
65
# These three are overridable in subclasses.
64
66
@@ -99,6 +101,8 @@ def _format(self):
99
101
result += f' _putters[{ len (self ._putters )} ]'
100
102
if self ._unfinished_tasks :
101
103
result += f' tasks={ self ._unfinished_tasks } '
104
+ if not self ._is_alive ():
105
+ result += f' state={ self ._shutdown_state .value } '
102
106
return result
103
107
104
108
def qsize (self ):
@@ -130,8 +134,10 @@ async def put(self, item):
130
134
131
135
Put an item into the queue. If the queue is full, wait until a free
132
136
slot is available before adding item.
137
+
138
+ Raises QueueShutDown if the queue has been shut down.
133
139
"""
134
- if self .shutdown_state != _queue_alive :
140
+ if not self ._is_alive () :
135
141
raise QueueShutDown
136
142
while self .full ():
137
143
putter = self ._get_loop ().create_future ()
@@ -145,23 +151,25 @@ async def put(self, item):
145
151
self ._putters .remove (putter )
146
152
except ValueError :
147
153
# The putter could be removed from self._putters by a
148
- # previous get_nowait call.
154
+ # previous get_nowait call or a shutdown call .
149
155
pass
150
156
if not self .full () and not putter .cancelled ():
151
157
# We were woken up by get_nowait(), but can't take
152
158
# the call. Wake up the next in line.
153
159
self ._wakeup_next (self ._putters )
154
160
raise
155
- if self .shutdown_state != _queue_alive :
161
+ if not self ._is_alive () :
156
162
raise QueueShutDown
157
163
return self .put_nowait (item )
158
164
159
165
def put_nowait (self , item ):
160
166
"""Put an item into the queue without blocking.
161
167
162
168
If no free slot is immediately available, raise QueueFull.
169
+
170
+ Raises QueueShutDown if the queue has been shut down.
163
171
"""
164
- if self .shutdown_state != _queue_alive :
172
+ if not self ._is_alive () :
165
173
raise QueueShutDown
166
174
if self .full ():
167
175
raise QueueFull
@@ -174,11 +182,14 @@ async def get(self):
174
182
"""Remove and return an item from the queue.
175
183
176
184
If queue is empty, wait until an item is available.
185
+
186
+ Raises QueueShutDown if the queue has been shut down and is empty, or
187
+ if the queue has been shut down immediately.
177
188
"""
178
- if self .shutdown_state == _queue_shutdown_immediate :
189
+ if self ._is_shutdown_immediate () :
179
190
raise QueueShutDown
180
191
while self .empty ():
181
- if self .shutdown_state != _queue_alive :
192
+ if self ._is_shutdown () :
182
193
raise QueueShutDown
183
194
getter = self ._get_loop ().create_future ()
184
195
self ._getters .append (getter )
@@ -191,28 +202,32 @@ async def get(self):
191
202
self ._getters .remove (getter )
192
203
except ValueError :
193
204
# The getter could be removed from self._getters by a
194
- # previous put_nowait call.
205
+ # previous put_nowait call,
206
+ # or a shutdown call.
195
207
pass
196
208
if not self .empty () and not getter .cancelled ():
197
209
# We were woken up by put_nowait(), but can't take
198
210
# the call. Wake up the next in line.
199
211
self ._wakeup_next (self ._getters )
200
212
raise
201
- if self .shutdown_state == _queue_shutdown_immediate :
213
+ if self ._is_shutdown_immediate () :
202
214
raise QueueShutDown
203
215
return self .get_nowait ()
204
216
205
217
def get_nowait (self ):
206
218
"""Remove and return an item from the queue.
207
219
208
220
Return an item if one is immediately available, else raise QueueEmpty.
221
+
222
+ Raises QueueShutDown if the queue has been shut down and is empty, or
223
+ if the queue has been shut down immediately.
209
224
"""
225
+ if self ._is_shutdown_immediate ():
226
+ raise QueueShutDown
210
227
if self .empty ():
211
- if self .shutdown_state != _queue_alive :
228
+ if self ._is_shutdown () :
212
229
raise QueueShutDown
213
230
raise QueueEmpty
214
- elif self .shutdown_state == _queue_shutdown_immediate :
215
- raise QueueShutDown
216
231
item = self ._get ()
217
232
self ._wakeup_next (self ._putters )
218
233
return item
@@ -230,7 +245,11 @@ def task_done(self):
230
245
231
246
Raises ValueError if called more times than there were items placed in
232
247
the queue.
248
+
249
+ Raises QueueShutDown if the queue has been shut down immediately.
233
250
"""
251
+ if self ._is_shutdown_immediate ():
252
+ raise QueueShutDown
234
253
if self ._unfinished_tasks <= 0 :
235
254
raise ValueError ('task_done() called too many times' )
236
255
self ._unfinished_tasks -= 1
@@ -244,9 +263,15 @@ async def join(self):
244
263
queue. The count goes down whenever a consumer calls task_done() to
245
264
indicate that the item was retrieved and all work on it is complete.
246
265
When the count of unfinished tasks drops to zero, join() unblocks.
266
+
267
+ Raises QueueShutDown if the queue has been shut down immediately.
247
268
"""
269
+ if self ._is_shutdown_immediate ():
270
+ raise QueueShutDown
248
271
if self ._unfinished_tasks > 0 :
249
272
await self ._finished .wait ()
273
+ if self ._is_shutdown_immediate ():
274
+ raise QueueShutDown
250
275
251
276
def shutdown (self , immediate = False ):
252
277
"""Shut-down the queue, making queue gets and puts raise.
@@ -257,19 +282,40 @@ def shutdown(self, immediate=False):
257
282
All blocked callers of put() will be unblocked, and also get()
258
283
and join() if 'immediate'. The QueueShutDown exception is raised.
259
284
"""
285
+ if self ._is_shutdown_immediate ():
286
+ return
287
+ # here _shutdown_state is ALIVE or SHUTDOWN
260
288
if immediate :
261
- self .shutdown_state = _queue_shutdown_immediate
289
+ self ._set_shutdown_immediate ()
262
290
while self ._getters :
263
291
getter = self ._getters .popleft ()
264
292
if not getter .done ():
265
293
getter .set_result (None )
294
+ # Release all 'blocked' tasks/coros in `join()`
295
+ self ._finished .set ()
266
296
else :
267
- self .shutdown_state = _queue_shutdown
297
+ self ._set_shutdown ()
268
298
while self ._putters :
269
299
putter = self ._putters .popleft ()
270
300
if not putter .done ():
271
301
putter .set_result (None )
272
302
303
+ def _is_alive (self ):
304
+ return self ._shutdown_state is _QueueState .ALIVE
305
+
306
+ def _is_shutdown (self ):
307
+ return self ._shutdown_state is _QueueState .SHUTDOWN
308
+
309
+ def _is_shutdown_immediate (self ):
310
+ return self ._shutdown_state is _QueueState .SHUTDOWN_IMMEDIATE
311
+
312
+ def _set_shutdown (self ):
313
+ self ._shutdown_state = _QueueState .SHUTDOWN
314
+
315
+ def _set_shutdown_immediate (self ):
316
+ self ._shutdown_state = _QueueState .SHUTDOWN_IMMEDIATE
317
+
318
+
273
319
class PriorityQueue (Queue ):
274
320
"""A subclass of Queue; retrieves entries in priority order (lowest first).
275
321
0 commit comments