from queue import Queue | |
import threading | |
class GeneratedQueue(Queue): | |
""" | |
usage: | |
for item in q.each(): | |
do_work(item) | |
accepts both generators and iterables | |
""" | |
def __init__(self, generator, preload=0): | |
super(GeneratedQueue, self).__init__() | |
self.lock = threading.Lock() | |
try: | |
self.generator = iter(generator) | |
except TypeError: | |
self.generator = generator() | |
for i in range(preload + 1): | |
self.load_next() | |
def load_next(self): | |
try: | |
with self.lock: | |
self.put(next(self.generator)) | |
except StopIteration: | |
pass | |
def each(self): | |
if self.empty(): | |
self.load_next() | |
while not self.empty(): | |
yield self.get() | |
self.task_done() | |
self.load_next() | |
def do_work(item): | |
print('too EZ', item) | |
def worker(): | |
for item in q.each(): | |
do_work(item) | |
print('stop') | |
def generator(): | |
for i in range(11): | |
yield i | |
q = GeneratedQueue(generator) | |
num_worker_threads = 4 | |
for i in range(num_worker_threads): | |
t = threading.Thread(target=worker) | |
t.daemon = True | |
t.start() | |
q.join() | |
print('finish') | |
q = GeneratedQueue([1,2,4]) | |
num_worker_threads = 4 | |
for i in range(num_worker_threads): | |
t = Thread(target=worker) | |
t.daemon = True | |
t.start() | |
q.join() | |
print('finish') |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment