Support using asyncio async generators in trio.

Author: miracle2k

trio_asyncio/adapter.py

@@ -1,6 +1,7 @@
 # This code implements a clone of the asyncio mainloop which hooks into
 # Trio.
 
+import inspect
 import trio_asyncio
 
 # import logging
@@ -12,9 +13,15 @@
 
 
 def trio2aio(proc):
-    @wraps(proc)
-    async def call(*args):
-        return await trio_asyncio.run_asyncio(proc, *args)
+    if inspect.isasyncgenfunction(proc):
+        @wraps(proc)
+        def call(*args):
+            return trio_asyncio.wrap_generator(proc, *args)
+
+    else:
+        @wraps(proc)
+        async def call(*args):
+            return await trio_asyncio.run_asyncio(proc, *args)
 
     return call
 

trio_asyncio/base.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import inspect
 import math
 import trio
 import heapq
@@ -11,7 +12,7 @@
 
 from selectors import _BaseSelectorImpl, EVENT_READ, EVENT_WRITE
 
-from .util import run_future
+from .util import run_future, run_generator
 
 try:
     from trio.hazmat import wait_for_child
@@ -212,6 +213,11 @@ async def run_coroutine(self, coro):
         coro = asyncio.ensure_future(coro, loop=self)
         return await run_future(coro)
 
+    def wrap_generator(self, gen):
+        # if inspect.isasyncgen(f):
+        # return self.wrap_generator(f)
+        return run_generator(self, gen())
+
     async def run_asyncio(self, proc, *args):
         """Run an asyncio function or method from Trio.
 

trio_asyncio/loop.py

@@ -25,6 +25,7 @@
     'run_future',
     'run_coroutine',
     'run_asyncio',
+    'wrap_generator',
     'TrioChildWatcher',
     'TrioPolicy',
 ]
@@ -169,6 +170,13 @@ def __exit__(self, *tb):
         self.close()
 
 
+def wrap_generator(proc, *args):
+    loop = asyncio.get_event_loop()
+    if not isinstance(loop, TrioEventLoop):
+        raise RuntimeError("Need to run in a trio_asyncio.open_loop() context")
+    return loop.wrap_generator(proc, *args)
+
+
 async def run_asyncio(proc, *args):
     loop = asyncio.get_event_loop()
     if not isinstance(loop, TrioEventLoop):

trio_asyncio/util.py

@@ -45,3 +45,48 @@ def abort_cb(raise_cancel_arg):
                 sys.exc_info()[1].__cause__ = exc
         else:
             raise
+
+
+STOP = object()
+
+async def run_generator(loop, async_generator):
+    task = trio.hazmat.current_task()
+    raise_cancel = None
+
+    async def consume_next():
+        try:
+            item = await async_generator.__anext__()
+        except StopAsyncIteration:
+            item = STOP
+
+        trio.hazmat.reschedule(task, trio.hazmat.Value(value=item))
+        #trio.hazmat.reschedule(task, STOP)
+
+    def abort_cb(raise_cancel_arg):
+        # Save the cancel-raising function
+        nonlocal raise_cancel
+        raise_cancel = raise_cancel_arg
+        # XXX: we need to cancel any actice consume_next() call.
+        # Keep waiting
+        return trio.hazmat.Abort.FAILED
+
+    try:
+        while True:
+            # schedule that we read the next one from the iterator
+            asyncio.ensure_future(consume_next(), loop=loop)
+
+            item = await trio.hazmat.wait_task_rescheduled(abort_cb)
+            if item == STOP:
+                break
+            yield item
+
+    except asyncio.CancelledError as exc:
+        if raise_cancel is not None:
+            try:
+                raise_cancel()
+            finally:
+                # Try to preserve the exception chain,
+                # for more detailed tracebacks
+                sys.exc_info()[1].__cause__ = exc
+        else:
+            raise