You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
51 lines
1.6 KiB
51 lines
1.6 KiB
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
|