From a9985c0c01de7fdbe06293536e4c3e18102f9ac4 Mon Sep 17 00:00:00 2001
From: Tyler Goodlet <jgbt@protonmail.com>
Date: Fri, 3 Jul 2020 17:33:46 -0400
Subject: [PATCH] Wow, fix all the broken async func invoking code..

Clearly this wasn't developed against a task that spawned just an async
func in `asyncio`.. Fix all that and remove a bunch of unnecessary func
layers. Add provisional support for the target receiving the `to_trio`
and `from_trio` channels and for the @tractor.stream marker.
---
 tractor/to_asyncio.py | 67 +++++++++++++++++++++++--------------------
 1 file changed, 36 insertions(+), 31 deletions(-)

diff --git a/tractor/to_asyncio.py b/tractor/to_asyncio.py
index a45904a..fb8f4cd 100644
--- a/tractor/to_asyncio.py
+++ b/tractor/to_asyncio.py
@@ -4,7 +4,6 @@ Infection apis for ``asyncio`` loops running ``trio`` using guest mode.
 import asyncio
 import inspect
 from typing import (
-    Any,
     Callable,
     AsyncGenerator,
     Awaitable,
@@ -19,47 +18,55 @@ from ._state import current_actor
 log = get_logger(__name__)
 
 
-__all__ = ['run_task']
+__all__ = ['run_task', 'run_as_asyncio_guest']
 
 
 async def _invoke(
     from_trio: trio.abc.ReceiveChannel,
     to_trio: asyncio.Queue,
     coro: Awaitable,
-) -> Union[AsyncGenerator, Awaitable]:
-    """Await or stream awaiable object based on type into
+) -> None:
+    """Await or stream awaiable object based on ``coro`` type into
     ``trio`` memory channel.
+
+    ``from_trio`` might eventually be used here for bidirectional streaming.
     """
-    async def stream_from_gen(c):
-        async for item in c:
-            to_trio.send_nowait(item)
-
-    async def just_return(c):
-        to_trio.send_nowait(await c)
-
     if inspect.isasyncgen(coro):
-        return await stream_from_gen(coro)
+        async for item in coro:
+            to_trio.send_nowait(item)
     elif inspect.iscoroutine(coro):
-        return await coro
+        to_trio.send_nowait(await coro)
 
 
 async def run_task(
     func: Callable,
+    *,
     qsize: int = 2**10,
+    _treat_as_stream: bool = False,
     **kwargs,
-) -> Any:
+) -> Union[AsyncGenerator, Awaitable]:
     """Run an ``asyncio`` async function or generator in a task, return
     or stream the result back to ``trio``.
     """
-    assert current_actor()._infected_aio
+    assert current_actor().is_infected_aio()
 
     # ITC (inter task comms)
     from_trio = asyncio.Queue(qsize)
     to_trio, from_aio = trio.open_memory_channel(qsize)
 
-    # allow target func to accept/stream results manually
-    kwargs['to_trio'] = to_trio
-    kwargs['from_trio'] = to_trio
+    args = tuple(inspect.getfullargspec(func).args)
+
+    if getattr(func, '_tractor_steam_function', None):
+        # the assumption is that the target async routine accepts the
+        # send channel then it intends to yield more then one return
+        # value otherwise it would just return ;P
+        _treat_as_stream = True
+
+    # allow target func to accept/stream results manually by name
+    if 'to_trio' in args:
+        kwargs['to_trio'] = to_trio
+    if 'from_trio' in args:
+        kwargs['from_trio'] = to_trio
 
     coro = func(**kwargs)
 
@@ -70,7 +77,6 @@ async def run_task(
     task = asyncio.create_task(_invoke(from_trio, to_trio, coro))
     err = None
 
-    # XXX: I'm not sure this actually does anything...
     def cancel_trio(task):
         """Cancel the calling ``trio`` task on error.
         """
@@ -80,17 +86,8 @@ async def run_task(
 
     task.add_done_callback(cancel_trio)
 
-    # determine return type async func vs. gen
-    if inspect.isasyncgen(coro):
-        # simple async func
-        async def result():
-            with cancel_scope:
-                return await from_aio.get()
-            if cancel_scope.cancelled_caught and err:
-                raise err
-
-    elif inspect.iscoroutine(coro):
-        # asycn gen
+    # asycn gen
+    if inspect.isasyncgen(coro) or _treat_as_stream:
         async def result():
             with cancel_scope:
                 async with from_aio:
@@ -99,7 +96,15 @@ async def run_task(
             if cancel_scope.cancelled_caught and err:
                 raise err
 
-    return result()
+        return result()
+
+    # simple async func
+    elif inspect.iscoroutine(coro):
+        with cancel_scope:
+            result = await from_aio.receive()
+            return result
+        if cancel_scope.cancelled_caught and err:
+            raise err
 
 
 def run_as_asyncio_guest(