|
|
- import asyncio
- import threading
- import tornado.queues
-
- # As long as we support Python35, we use this library to get as async
- # generators: https://pypi.org/project/async_generator/
- from async_generator import async_generator, yield_
-
-
- class ThreadedAsyncGenerator(threading.Thread):
- def __init__(self, main_ioloop, fn, *args, **kwargs):
- super(ThreadedAsyncGenerator, self).__init__()
- self.main_ioloop = main_ioloop
- self.fn = fn
- self.args = args
- self.kwargs = kwargs
- self.queue = tornado.queues.Queue()
- self.start()
-
- def run(self):
- ioloop_in_thread = asyncio.new_event_loop()
- asyncio.set_event_loop(ioloop_in_thread)
- return ioloop_in_thread.run_until_complete(self._run())
-
- async def _run(self):
- async for item in self.fn(*self.args, **self.kwargs):
- def thread_safe_put(item=item):
- self.queue.put(item)
- self.main_ioloop.call_soon_threadsafe(thread_safe_put)
-
- def thread_safe_end():
- self.queue.put(StopIteration)
- self.main_ioloop.call_soon_threadsafe(thread_safe_end)
-
- @async_generator
- async def __aiter__(self):
- while True:
- value = await self.queue.get()
- if value == StopIteration:
- break
- await yield_(value)
-
-
- def async_generator_to_thread(fn):
- """Calls an async generator function fn in a thread and async returns the results"""
- ioloop = asyncio.get_event_loop()
-
- def wrapper(*args, **kwargs):
- gen = ThreadedAsyncGenerator(ioloop, fn, *args, **kwargs)
- return gen
- return wrapper
|