Add the ability for @async_generator to produce a native async generator

async_generator/__init__.py

@@ -7,11 +7,18 @@
     isasyncgenfunction,
     get_asyncgen_hooks,
     set_asyncgen_hooks,
+    supports_native_asyncgens,
 )
 from ._util import aclosing, asynccontextmanager
+from functools import partial as _partial
+
+async_generator_native = _partial(async_generator, allow_native=True)
+async_generator_legacy = _partial(async_generator, allow_native=False)
 
 __all__ = [
     "async_generator",
+    "async_generator_native",
+    "async_generator_legacy",
     "yield_",
     "yield_from_",
     "aclosing",
@@ -20,4 +27,5 @@
     "asynccontextmanager",
     "get_asyncgen_hooks",
     "set_asyncgen_hooks",
+    "supports_native_asyncgens",
 ]

async_generator/_impl.py

@@ -1,6 +1,6 @@
 import sys
-from functools import wraps
-from types import coroutine, CodeType
+from functools import wraps, partial
+from types import coroutine, CodeType, FunctionType
 import inspect
 from inspect import (
     getcoroutinestate, CORO_CREATED, CORO_CLOSED, CORO_SUSPENDED
@@ -50,8 +50,11 @@ def inner():
 
 if sys.implementation.name == "cpython" and sys.version_info >= (3, 6):
     # On 3.6, with native async generators, we want to use the same
-    # wrapper type that native generators use. This lets @async_generators
-    # yield_from_ native async generators and vice versa.
+    # wrapper type that native generators use. This lets @async_generator
+    # create a native async generator under most circumstances, which
+    # improves performance while still permitting 3.5-compatible syntax.
+    # It also lets non-native @async_generators (e.g. those that return
+    # non-None values) yield_from_ native ones and vice versa.
 
     import ctypes
     from types import AsyncGeneratorType, GeneratorType
@@ -315,10 +318,11 @@ def set_asyncgen_hooks(firstiter=UNSPECIFIED, finalizer=UNSPECIFIED):
 
 
 class AsyncGenerator:
-    def __init__(self, coroutine):
+    def __init__(self, coroutine, *, warn_on_native_differences=False):
         self._coroutine = coroutine
+        self._warn_on_native_differences = warn_on_native_differences
         self._it = coroutine.__await__()
-        self.ag_running = False
+        self._running = False
         self._finalizer = None
         self._closed = False
         self._hooks_inited = False
@@ -351,6 +355,28 @@ def ag_code(self):
     def ag_frame(self):
         return self._coroutine.cr_frame
 
+    @property
+    def ag_await(self):
+        return self._coroutine.cr_await
+
+    @property
+    def ag_running(self):
+        if self._running != self._coroutine.cr_running and self._warn_on_native_differences:
+            import warnings
+            warnings.warn(
+                "Native async generators incorrectly set ag_running = False "
+                "when the generator is awaiting a trap to the event loop and "
+                "not suspended via a yield to its caller. Your code examines "
+                "ag_running under such conditions, and will change behavior "
+                "when async_generator starts using native generators by default "
+                "(where available) in the next release. "
+                "Use @async_generator_legacy to keep the current behavior, or "
+                "@async_generator_native if you're OK with the change.",
+                category=FutureWarning,
+                stacklevel=2
+            )
+        return self._running
+
     ################################################################
     # Core functionality
     ################################################################
@@ -387,10 +413,10 @@ async def step():
             if self.ag_running:
                 raise ValueError("async generator already executing")
             try:
-                self.ag_running = True
+                self._running = True
                 return await ANextIter(self._it, start_fn, *args)
             finally:
-                self.ag_running = False
+                self._running = False
 
         return step()
 
@@ -462,10 +488,128 @@ def __del__(self):
     collections.abc.AsyncGenerator.register(AsyncGenerator)
 
 
-def async_generator(coroutine_maker):
-    @wraps(coroutine_maker)
+def _find_return_of_not_none(co):
+    """Inspect the code object *co* for the presence of return statements that
+    might return a value other than None. If any such statement is found,
+    return the source line number (in file ``co.co_filename``) on which it occurs.
+    If all return statements definitely return the value None, return None.
+    """
+
+    # 'return X' for simple/constant X seems to always compile to
+    # LOAD_CONST + RETURN_VALUE immediately following one another,
+    # and LOAD_CONST(None) + RETURN_VALUE definitely does mean return None,
+    # so we'll search for RETURN_VALUE not preceded by LOAD_CONST or preceded
+    # by a LOAD_CONST that does not load None.
+    import dis
+    current_line = co.co_firstlineno
+    prev_inst = None
+    for inst in dis.Bytecode(co):
+        if inst.starts_line is not None:
+            current_line = inst.starts_line
+        if inst.opname == "RETURN_VALUE" and (prev_inst is None or
+                                              prev_inst.opname != "LOAD_CONST"
+                                              or prev_inst.argval is not None):
+            return current_line
+        prev_inst = inst
+    return None
+
+
+def _as_native_asyncgen_function(afn):
+    """Given a non-generator async function *afn*, which contains some ``await yield_(...)``
+    and/or ``await yield_from_(...)`` calls, create the analogous async generator function.
+    *afn* must not return values other than None; doing so is likely to crash the interpreter.
+    Use :func:`_find_return_of_not_none` to check this.
+    """
+
+    # An async function that contains 'await yield_()' statements is a perfectly
+    # cromulent async generator, except that it doesn't get marked with CO_ASYNC_GENERATOR
+    # because it doesn't contain any 'yield' statements. Create a new code object that
+    # does have CO_ASYNC_GENERATOR set. This is preferable to using a wrapper because
+    # it makes ag_code and ag_frame work the same way they would for a native async
+    # generator with 'yield' statements, and the same as for an async_generator
+    # asyncgen with allow_native=False.
+
+    from inspect import CO_COROUTINE, CO_ASYNC_GENERATOR
+    co = afn.__code__
+    new_code = CodeType(
+        co.co_argcount, co.co_kwonlyargcount, co.co_nlocals, co.co_stacksize,
+        (co.co_flags & ~CO_COROUTINE) | CO_ASYNC_GENERATOR, co.co_code,
+        co.co_consts, co.co_names, co.co_varnames, co.co_filename, co.co_name,
+        co.co_firstlineno, co.co_lnotab, co.co_freevars, co.co_cellvars
+    )
+    asyncgen_fn = FunctionType(
+        new_code, afn.__globals__, afn.__name__, afn.__defaults__,
+        afn.__closure__
+    )
+    asyncgen_fn.__kwdefaults__ = afn.__kwdefaults__
+    return wraps(afn)(asyncgen_fn)
+
+
+def async_generator(afn=None, *, allow_native=None, uses_return=None):
+    if afn is None:
+        return partial(
+            async_generator,
+            uses_return=uses_return,
+            allow_native=allow_native
+        )
+
+    uses_wrapper = False
+    if not inspect.iscoroutinefunction(afn):
+        underlying = afn
+        while hasattr(underlying, "__wrapped__"):
+            underlying = getattr(underlying, "__wrapped__")
+            if inspect.iscoroutinefunction(underlying):
+                break
+        else:
+            raise TypeError(
+                "expected an async function, not {!r}".format(
+                    type(afn).__name__
+                )
+            )
+        # A sync wrapper around an async function is fine, in the sense
+        # that we can call it to get a coroutine just like we could for
+        # an async function; but it's a bit suboptimal, in the sense that
+        # we can't turn it into an async generator. One way to get here
+        # is to put trio's @enable_ki_protection decorator below
+        # @async_generator rather than above it.
+        uses_wrapper = True
+
+    if sys.implementation.name == "cpython" and not uses_wrapper:
+        # 'return' statements with non-None arguments are syntactically forbidden when
+        # compiling a true async generator, but the flags mutation in
+        # _convert_to_native_asyncgen_function sidesteps that check, which could raise
+        # an assertion in genobject.c when a non-None return is executed.
+        # To avoid crashing the interpreter due to a user error, we need to examine the
+        # code of the function we're converting.
+
+        co = afn.__code__
+        nontrivial_return_line = _find_return_of_not_none(co)
+        seems_to_use_return = nontrivial_return_line is not None
+        if uses_return is None:
+            uses_return = seems_to_use_return
+        elif uses_return != seems_to_use_return:
+            prefix = "{} declared using @async_generator(uses_return={}) but ".format(
+                afn.__qualname__, uses_return
+            )
+            if seems_to_use_return:
+                raise RuntimeError(
+                    prefix + "might return a value other than None at {}:{}".
+                    format(co.co_filename, nontrivial_return_line)
+                )
+            else:
+                raise RuntimeError(
+                    prefix + "never returns a value other than None"
+                )
+
+        if allow_native and not uses_return and supports_native_asyncgens:
+            return _as_native_asyncgen_function(afn)
+
+    @wraps(afn)
     def async_generator_maker(*args, **kwargs):
-        return AsyncGenerator(coroutine_maker(*args, **kwargs))
+        return AsyncGenerator(
+            afn(*args, **kwargs),
+            warn_on_native_differences=(allow_native is not False)
+        )
 
     async_generator_maker._async_gen_function = id(async_generator_maker)
     return async_generator_maker

async_generator/_tests/conftest.py

@@ -34,3 +34,41 @@ def wrapper(**kwargs):
                 pass
 
         pyfuncitem.obj = wrapper
+
+
+# On CPython 3.6+, tests that accept an 'async_generator' fixture
+# will be called twice, once to test behavior with legacy pure-Python
+# async generators (i.e. _impl.AsyncGenerator) and once to test behavior
+# with wrapped native async generators. Tests should decorate their
+# asyncgens with their local @async_generator, and may inspect
+# async_generator.is_native to know which sort is being used.
+# On CPython 3.5 or PyPy, where native async generators do not
+# exist, async_generator will always call _impl.async_generator()
+# and tests will be invoked only once.
+
+from .. import supports_native_asyncgens
+maybe_native = pytest.param(
+    "native",
+    marks=pytest.mark.skipif(
+        not supports_native_asyncgens,
+        reason="native async generators are not supported on this Python version"
+    )
+)
+
+
+@pytest.fixture(params=["legacy", maybe_native])
+def async_generator(request):
+    from .. import async_generator as real_async_generator
+
+    def wrapper(afn=None, *, uses_return=False):
+        if afn is None:
+            return partial(wrapper, uses_return=uses_return)
+        if request.param == "native":
+            return real_async_generator(afn, allow_native=True)
+        # make sure the uses_return= argument matches what async_generator detects
+        real_async_generator(afn, allow_native=False, uses_return=uses_return)
+        # make sure autodetection without uses_return= works too
+        return real_async_generator(afn, allow_native=False)
+
+    wrapper.is_native = request.param == "native"
+    return wrapper