forked from cool-RR/python_toolbox
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathqueue_tools.py
More file actions
104 lines (81 loc) · 2.98 KB
/
queue_tools.py
File metadata and controls
104 lines (81 loc) · 2.98 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
# Copyright 2009-2017 Ram Rachum.
# This program is distributed under the MIT license.
'''Defines various functions for working with queues.'''
from __future__ import generator_stop
import queue as queue_module
import sys
from python_toolbox import caching
from python_toolbox import import_tools
def is_multiprocessing_queue(queue):
'''Return whether `queue` is a multiprocessing queue.'''
return queue.__module__.startswith('multiprocessing')
def dump(queue):
'''
Empty all pending items in a queue and return them in a list.
Use only when no other processes/threads are reading from the queue.
'''
return list(iterate(queue))
def iterate(queue, block=False, limit_to_original_size=False,
_prefetch_if_no_qsize=False):
'''
Iterate over the items in the queue.
`limit_to_original_size=True` will limit the number of the items fetched to
the original number of items in the queue in the beginning.
'''
if limit_to_original_size:
if is_multiprocessing_queue(queue) and \
not _platform_supports_multiprocessing_qsize():
if _prefetch_if_no_qsize:
yield from dump(queue)
return
raise NotImplementedError(
"This platform doesn't support `qsize` for `multiprocessing` "
"queues, so you can't iterate on it while limiting to "
"original queue size. What you can do is set "
"`_prefetch_if_no_qsize=True` to have the entire queue "
"prefetched before yielding the items."
)
for _ in range(queue.qsize()):
try:
yield queue.get(block=block)
except queue_module.Empty:
return
else: # not limit_to_original_size
while True:
try:
yield queue.get(block=block)
except queue_module.Empty:
return
def get_item(queue, i):
'''
Get an item from the queue by index number without removing any items.
Note: This was designed for `Queue.Queue`. Don't try to use this, for
example, on `multiprocessing.Queue`.
'''
with queue.mutex:
return queue.queue[i]
def queue_as_list(queue):
'''
Get all the items in the queue as a `list` without removing them.
Note: This was designed for `Queue.Queue`. Don't try to use this, for
example, on `multiprocessing.Queue`.
'''
with queue.mutex:
return list(queue.queue)
@caching.cache()
def _platform_supports_multiprocessing_qsize():
'''
Return whether this platform supports `multiprocessing.Queue().qsize()`.
I'm looking at you, Mac OS.
'''
if 'multiprocessing' not in sys.modules:
if not import_tools.exists('multiprocessing'):
return False
import multiprocessing
multiprocessing_queue = multiprocessing.Queue()
try:
multiprocessing_queue.qsize()
except NotImplementedError:
return False
else:
return True