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

4 years ago
  1. import asyncio
  2. import threading
  3. import tornado.queues
  4. # As long as we support Python35, we use this library to get as async
  5. # generators: https://pypi.org/project/async_generator/
  6. from async_generator import async_generator, yield_
  7. class ThreadedAsyncGenerator(threading.Thread):
  8. def __init__(self, main_ioloop, fn, *args, **kwargs):
  9. super(ThreadedAsyncGenerator, self).__init__()
  10. self.main_ioloop = main_ioloop
  11. self.fn = fn
  12. self.args = args
  13. self.kwargs = kwargs
  14. self.queue = tornado.queues.Queue()
  15. self.start()
  16. def run(self):
  17. ioloop_in_thread = asyncio.new_event_loop()
  18. asyncio.set_event_loop(ioloop_in_thread)
  19. return ioloop_in_thread.run_until_complete(self._run())
  20. async def _run(self):
  21. async for item in self.fn(*self.args, **self.kwargs):
  22. def thread_safe_put(item=item):
  23. self.queue.put(item)
  24. self.main_ioloop.call_soon_threadsafe(thread_safe_put)
  25. def thread_safe_end():
  26. self.queue.put(StopIteration)
  27. self.main_ioloop.call_soon_threadsafe(thread_safe_end)
  28. @async_generator
  29. async def __aiter__(self):
  30. while True:
  31. value = await self.queue.get()
  32. if value == StopIteration:
  33. break
  34. await yield_(value)
  35. def async_generator_to_thread(fn):
  36. """Calls an async generator function fn in a thread and async returns the results"""
  37. ioloop = asyncio.get_event_loop()
  38. def wrapper(*args, **kwargs):
  39. gen = ThreadedAsyncGenerator(ioloop, fn, *args, **kwargs)
  40. return gen
  41. return wrapper