Make yield_ and yield_from_ interoperate with native async generators again

async_generator/_impl.py

@@ -1,94 +1,138 @@
 import sys
 from functools import wraps
-from types import coroutine
+from types import coroutine, CodeType
 import inspect
 from inspect import (
     getcoroutinestate, CORO_CREATED, CORO_CLOSED, CORO_SUSPENDED
 )
 import collections.abc
 
+# An async generator object (whether native in 3.6+ or the pure-Python
+# version implemented below) is basically an async function with some
+# extra wrapping logic. As an async function, it can call other async
+# functions, which will probably at some point call a function that uses
+# 'yield' to send traps to the event loop. Async generators also need
+# to be able to send values to the context in which the generator is
+# being iterated, and it's awfully convenient to be able to do that
+# using 'yield' too. To distinguish between these two streams of
+# yielded objects, the traps intended for the event loop are yielded
+# as-is, and the values intended for the context that's iterating the
+# generator are wrapped in some wrapper object (YieldWrapper here, or
+# an internal Python type called AsyncGenWrappedValue in the native
+# async generator implementation) before being yielded.
+# The __anext__(), asend(), and athrow() methods of an async generator
+# iterate the underlying async function until a wrapped value is received,
+# and any non-wrapped values are passed through to the event loop.
+
+# These functions are syntactically valid only on 3.6+, so we conditionally
+# exec() the code defining them.
+_native_asyncgen_helpers = """
+async def _wrapper():
+    holder = [None]
+    while True:
+        # The simpler "value = None; while True: value = yield value"
+        # would hold a reference to the most recently wrapped value
+        # after it has been yielded out (until the next value to wrap
+        # comes in), so we use a one-element list instead.
+        holder.append((yield holder.pop()))
+_wrapper = _wrapper()
+
+async def _unwrapper():
+    @coroutine
+    def inner():
+        holder = [None]
+        while True:
+            holder.append((yield holder.pop()))
+    await inner()
+    yield None
+_unwrapper = _unwrapper()
+"""
+
+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.
+
+    import ctypes
+    from types import AsyncGeneratorType, GeneratorType
+    exec(_native_asyncgen_helpers)
+
+    # Transmute _wrapper to a regular generator object by modifying the
+    # ob_type field. The code object inside _wrapper will still think it's
+    # associated with an async generator, so it will yield out
+    # AsyncGenWrappedValues when it encounters a 'yield' statement;
+    # but the generator object will think it's a normal non-async
+    # generator, so it won't unwrap them. This way, we can obtain
+    # AsyncGenWrappedValues as normal manipulable Python objects.
+    #
+    # This sort of object type transmutation is categorically a Sketchy
+    # Thing To Do, because the functions associated with the new type
+    # (including tp_dealloc and so forth) will be operating on a
+    # structure whose in-memory layout matches that of the old type.
+    # In this case, it's OK, because async generator objects are just
+    # generator objects plus a few extra fields at the end; and these
+    # fields are two integers and a NULL-until-first-iteration object
+    # pointer ag_finalizer, so they don't hold any resources that need
+    # to be cleaned up. (We transmute the async generator to a regular
+    # generator before we first iterate it, so ag_finalizer stays NULL
+    # for the lifetime of the object, so failing to drop a reference
+    # to it during deallocation doesn't actually leak anything.)
+    # We have a unit test that verifies that __sizeof__() for generators
+    # and async generators continues to follow this pattern in future
+    # Python versions.
+
+    _type_p = ctypes.c_void_p.from_address(
+        id(_wrapper) + object().__sizeof__() - ctypes.sizeof(ctypes.c_void_p)
+    )
+    assert _type_p.value == id(AsyncGeneratorType)
+    _type_p.value = id(GeneratorType)
 
-class YieldWrapper:
-    def __init__(self, payload):
-        self.payload = payload
+    supports_native_asyncgens = True
 
+    # Now _wrapper.send(x) returns an AsyncGenWrappedValue of x.
+    # We have to initially send(None) since the generator was just constructed;
+    # we remember the type of the return value (which is AsyncGenWrappedValue(None))
+    # to help with _is_wrapped.
+    YieldWrapper = type(_wrapper.send(None))
 
-def _wrap(value):
-    return YieldWrapper(value)
+    # Advance _unwrapper to its first yield statement, for use by _unwrap().
+    _unwrapper.asend(None).send(None)
 
+    # Performance note: compared to the non-native-supporting implementation below,
+    # this _wrap() is about the same speed (434 +- 16 nsec here, 456 +- 24 nsec below)
+    # but this _unwrap() is much slower (1.17 usec vs 167 nsec). Since _unwrap is only
+    # needed on non-native generators, and we plan to have most @async_generators use
+    # native generators on 3.6+, this seems acceptable.
 
-def _is_wrapped(box):
-    return isinstance(box, YieldWrapper)
+    _wrap = _wrapper.send
 
+    def _is_wrapped(box):
+        return isinstance(box, YieldWrapper)
 
-def _unwrap(box):
-    return box.payload
+    def _unwrap(box):
+        try:
+            _unwrapper.asend(box).send(None)
+        except StopIteration as e:
+            return e.value
+        else:
+            raise TypeError("not wrapped")
+else:
+    supports_native_asyncgens = False
 
+    class YieldWrapper:
+        __slots__ = ("payload",)
 
-# This is the magic code that lets you use yield_ and yield_from_ with native
-# generators.
-#
-# The old version worked great on Linux and MacOS, but not on Windows, because
-# it depended on _PyAsyncGenValueWrapperNew. The new version segfaults
-# everywhere, and I'm not sure why -- probably my lack of understanding
-# of ctypes and refcounts.
-#
-# There are also some commented out tests that should be re-enabled if this is
-# fixed:
-#
-# if sys.version_info >= (3, 6):
-#     # Use the same box type that the interpreter uses internally. This allows
-#     # yield_ and (more importantly!) yield_from_ to work in built-in
-#     # generators.
-#     import ctypes  # mua ha ha.
-#
-#     # We used to call _PyAsyncGenValueWrapperNew to create and set up new
-#     # wrapper objects, but that symbol isn't available on Windows:
-#     #
-#     #   https://github.com/python-trio/async_generator/issues/5
-#     #
-#     # Fortunately, the type object is available, but it means we have to do
-#     # this the hard way.
-#
-#     # We don't actually need to access this, but we need to make a ctypes
-#     # structure so we can call addressof.
-#     class _ctypes_PyTypeObject(ctypes.Structure):
-#         pass
-#     _PyAsyncGenWrappedValue_Type_ptr = ctypes.addressof(
-#         _ctypes_PyTypeObject.in_dll(
-#             ctypes.pythonapi, "_PyAsyncGenWrappedValue_Type"))
-#     _PyObject_GC_New = ctypes.pythonapi._PyObject_GC_New
-#     _PyObject_GC_New.restype = ctypes.py_object
-#     _PyObject_GC_New.argtypes = (ctypes.c_void_p,)
-#
-#     _Py_IncRef = ctypes.pythonapi.Py_IncRef
-#     _Py_IncRef.restype = None
-#     _Py_IncRef.argtypes = (ctypes.py_object,)
-#
-#     class _ctypes_PyAsyncGenWrappedValue(ctypes.Structure):
-#         _fields_ = [
-#             ('PyObject_HEAD', ctypes.c_byte * object().__sizeof__()),
-#             ('agw_val', ctypes.py_object),
-#         ]
-#     def _wrap(value):
-#         box = _PyObject_GC_New(_PyAsyncGenWrappedValue_Type_ptr)
-#         raw = ctypes.cast(ctypes.c_void_p(id(box)),
-#                           ctypes.POINTER(_ctypes_PyAsyncGenWrappedValue))
-#         raw.contents.agw_val = value
-#         _Py_IncRef(value)
-#         return box
-#
-#     def _unwrap(box):
-#         assert _is_wrapped(box)
-#         raw = ctypes.cast(ctypes.c_void_p(id(box)),
-#                           ctypes.POINTER(_ctypes_PyAsyncGenWrappedValue))
-#         value = raw.contents.agw_val
-#         _Py_IncRef(value)
-#         return value
-#
-#     _PyAsyncGenWrappedValue_Type = type(_wrap(1))
-#     def _is_wrapped(box):
-#         return isinstance(box, _PyAsyncGenWrappedValue_Type)
+        def __init__(self, payload):
+            self.payload = payload
+
+    def _wrap(value):
+        return YieldWrapper(value)
+
+    def _is_wrapped(box):
+        return isinstance(box, YieldWrapper)
+
+    def _unwrap(box):
+        return box.payload
 
 
 # The magic @coroutine decorator is how you write the bottom level of

async_generator/_tests/test_async_generator.py

@@ -5,6 +5,7 @@
 import collections.abc
 from functools import wraps
 import gc
+import inspect
 
 from .conftest import mock_sleep
 from .. import (
@@ -370,13 +371,12 @@ async def native_async_range(count):
     for i in range(count):
         yield i
 
-# XX uncomment if/when we re-enable the ctypes hacks:
-# async def native_async_range_twice(count):
-#     # make sure yield_from_ works inside a native async generator
-#     await yield_from_(async_range(count))
-#     yield None
-#     # make sure we can yield_from_ a native async generator
-#     await yield_from_(native_async_range(count))
+async def native_async_range_twice(count, async_range):
+    # make sure yield_from_ works inside a native async generator
+    await yield_from_(async_range(count))
+    yield None
+    # make sure we can yield_from_ a native async generator
+    await yield_from_(native_async_range(count))
     """
     )
 
@@ -400,11 +400,16 @@ async def yield_from_native():
 
         assert await collect(yield_from_native()) == [0, 1, 2]
 
-    # XX uncomment if/when we re-enable the ctypes hacks:
-    # if sys.version_info >= (3, 6):
-    #     assert await collect(native_async_range_twice(3)) == [
-    #         0, 1, 2, None, 0, 1, 2,
-    #     ]
+    if sys.version_info >= (3, 6):
+        assert await collect(native_async_range_twice(3, async_range)) == [
+            0,
+            1,
+            2,
+            None,
+            0,
+            1,
+            2,
+        ]
 
 
 @async_generator
@@ -762,32 +767,47 @@ async def f():
 @pytest.mark.skipif(not hasattr(sys, "getrefcount"), reason="CPython only")
 def test_refcnt():
     x = object()
-    print(sys.getrefcount(x))
-    print(sys.getrefcount(x))
-    print(sys.getrefcount(x))
-    print(sys.getrefcount(x))
     base_count = sys.getrefcount(x)
     l = [_impl._wrap(x) for _ in range(100)]
-    print(sys.getrefcount(x))
-    print(sys.getrefcount(x))
-    print(sys.getrefcount(x))
-    assert sys.getrefcount(x) >= base_count + 100
+    assert sys.getrefcount(x) == base_count + 100
     l2 = [_impl._unwrap(box) for box in l]
-    assert sys.getrefcount(x) >= base_count + 200
-    print(sys.getrefcount(x))
-    print(sys.getrefcount(x))
-    print(sys.getrefcount(x))
-    print(sys.getrefcount(x))
+    assert sys.getrefcount(x) == base_count + 200
     del l
-    print(sys.getrefcount(x))
-    print(sys.getrefcount(x))
-    print(sys.getrefcount(x))
+    assert sys.getrefcount(x) == base_count + 100
     del l2
-    print(sys.getrefcount(x))
-    print(sys.getrefcount(x))
-    print(sys.getrefcount(x))
     assert sys.getrefcount(x) == base_count
-    print(sys.getrefcount(x))
+
+
+@pytest.mark.skipif(
+    not hasattr(None, "__sizeof__") or not hasattr(inspect, "isasyncgen"),
+    reason="CPython with native asyncgens only"
+)
+def test_gen_agen_size():
+    def gen():  # pragma: no cover
+        yield 42
+
+    dct = {}
+    exec("async def agen(): yield 50", dct)
+    agen = dct["agen"]
+
+    # As of CPython 3.7, an async generator object is a generator object plus one pointer
+    # (PyObject *ag_finalizer, nullptr before first iteration) and two ints
+    # (ag_hooks_inited, ag_closed). Since none of these members require any cleanup,
+    # our sketchy agen->gen type transmutation in _impl._wrapper is safe (at the time
+    # we do it, i.e., before first iteration). If the below assertion starts firing
+    # on a future Python version, someone will need to audit the new definition of
+    # PyAsyncGenObject (in CPython's Include/genobject.h) and make sure its new fields
+    # remain safe to ignore under the circumstances used in the _wrapper hack.
+    from ctypes import sizeof, c_int, c_void_p
+    gen_size = gen().__sizeof__()
+    agen_size = agen().__sizeof__()
+    expected_delta = 2 * sizeof(c_int) + sizeof(c_void_p)
+    assert gen_size + expected_delta == agen_size
+
+
+def test_unwrap_not_wrapped():
+    with pytest.raises((TypeError, AttributeError)):
+        _impl._unwrap(42)
 
 
 ################################################################

docs/source/reference.rst

@@ -28,6 +28,8 @@ Basically:
 That's it!
 
 
+.. _yieldfrom:
+
 Yield from
 ~~~~~~~~~~
 
@@ -47,9 +49,12 @@ But we do::
     async def wrap_load_json_lines(stream_reader):
         await yield_from_(load_json_lines(stream_reader))
 
-You can only use ``yield_from_`` inside an ``@async_generator``
-function, BUT the thing you PASS to ``yield_from_`` can be any kind of
-async iterator, including native async generators.
+You can use ``yield_from_`` inside a native async generator or an
+``@async_generator`` function, and the argument can be any async
+iterable. Remember that a native async generator is only created
+when an ``async def`` block contains at least one ``yield``
+statement; if you only ``yield_from_`` and never ``yield``,
+use the ``@async_generator`` decorator.
 
 Our ``yield_from_`` fully supports the classic ``yield from``
 semantics, including forwarding ``asend`` and ``athrow`` calls into

newsfragments/16.feature.rst

@@ -0,0 +1 @@
+On CPython 3.6+, restore the ability for a native async generator to ``yield_from_`` an ``@async_generator`` function and vice versa, by using the same value-wrapper type that Python uses internally if we're running on a Python version that supports native async generators. The new approach should be much less prone to refcounting bugs or to issues on different architectures.