import sys from functools import wraps from types import coroutine import inspect from inspect import ( getcoroutinestate, CORO_CREATED, CORO_CLOSED, CORO_SUSPENDED ) import collections.abc class YieldWrapper: 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 # 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) # The magic @coroutine decorator is how you write the bottom level of # coroutine stacks -- 'async def' can only use 'await' = yield from; but # eventually we must bottom out in a @coroutine that calls plain 'yield'. @coroutine def _yield_(value): return (yield _wrap(value)) # But we wrap the bare @coroutine version in an async def, because async def # has the magic feature that users can get warnings messages if they forget to # use 'await'. async def yield_(value=None): return await _yield_(value) async def yield_from_(delegate): # Transcribed with adaptations from: # # https://www.python.org/dev/peps/pep-0380/#formal-semantics # # This takes advantage of a sneaky trick: if an @async_generator-wrapped # function calls another async function (like yield_from_), and that # second async function calls yield_, then because of the hack we use to # implement yield_, the yield_ will actually propagate through yield_from_ # back to the @async_generator wrapper. So even though we're a regular # function, we can directly yield values out of the calling async # generator. def unpack_StopAsyncIteration(e): if e.args: return e.args[0] else: return None _i = type(delegate).__aiter__(delegate) if hasattr(_i, "__await__"): _i = await _i try: _y = await type(_i).__anext__(_i) except StopAsyncIteration as _e: _r = unpack_StopAsyncIteration(_e) else: while 1: try: _s = await yield_(_y) except GeneratorExit as _e: try: _m = _i.aclose except AttributeError: pass else: await _m() raise _e except BaseException as _e: _x = sys.exc_info() try: _m = _i.athrow except AttributeError: raise _e else: try: _y = await _m(*_x) except StopAsyncIteration as _e: _r = unpack_StopAsyncIteration(_e) break else: try: if _s is None: _y = await type(_i).__anext__(_i) else: _y = await _i.asend(_s) except StopAsyncIteration as _e: _r = unpack_StopAsyncIteration(_e) break return _r # This is the awaitable / iterator that implements asynciter.__anext__() and # friends. # # Note: we can be sloppy about the distinction between # # type(self._it).__next__(self._it) # # and # # self._it.__next__() # # because we happen to know that self._it is not a general iterator object, # but specifically a coroutine iterator object where these are equivalent. class ANextIter: def __init__(self, it, first_fn, *first_args): self._it = it self._first_fn = first_fn self._first_args = first_args def __await__(self): return self def __next__(self): if self._first_fn is not None: first_fn = self._first_fn first_args = self._first_args self._first_fn = self._first_args = None return self._invoke(first_fn, *first_args) else: return self._invoke(self._it.__next__) def send(self, value): return self._invoke(self._it.send, value) def throw(self, type, value=None, traceback=None): return self._invoke(self._it.throw, type, value, traceback) def _invoke(self, fn, *args): try: result = fn(*args) except StopIteration as e: # The underlying generator returned, so we should signal the end # of iteration. raise StopAsyncIteration(e.value) except StopAsyncIteration as e: # PEP 479 says: if a generator raises Stop(Async)Iteration, then # it should be wrapped into a RuntimeError. Python automatically # enforces this for StopIteration; for StopAsyncIteration we need # to it ourselves. raise RuntimeError( "async_generator raise StopAsyncIteration" ) from e if _is_wrapped(result): raise StopIteration(_unwrap(result)) else: return result UNSPECIFIED = object() try: from sys import get_asyncgen_hooks, set_asyncgen_hooks except ImportError: import threading asyncgen_hooks = collections.namedtuple( "asyncgen_hooks", ("firstiter", "finalizer") ) class _hooks_storage(threading.local): def __init__(self): self.firstiter = None self.finalizer = None _hooks = _hooks_storage() def get_asyncgen_hooks(): return asyncgen_hooks( firstiter=_hooks.firstiter, finalizer=_hooks.finalizer ) def set_asyncgen_hooks(firstiter=UNSPECIFIED, finalizer=UNSPECIFIED): if firstiter is not UNSPECIFIED: if firstiter is None or callable(firstiter): _hooks.firstiter = firstiter else: raise TypeError( "callable firstiter expected, got {}".format( type(firstiter).__name__ ) ) if finalizer is not UNSPECIFIED: if finalizer is None or callable(finalizer): _hooks.finalizer = finalizer else: raise TypeError( "callable finalizer expected, got {}".format( type(finalizer).__name__ ) ) class AsyncGenerator: # https://bitbucket.org/pypy/pypy/issues/2786: # PyPy implements 'await' in a way that requires the frame object # used to execute a coroutine to keep a weakref to that coroutine. # During a GC pass, weakrefs to all doomed objects are broken # before any of the doomed objects' finalizers are invoked. # If an AsyncGenerator is unreachable, its _coroutine probably # is too, and the weakref from ag._coroutine.cr_frame to # ag._coroutine will be broken before ag.__del__ can do its # one-turn close attempt or can schedule a full aclose() using # the registered finalization hook. It doesn't look like the # underlying issue is likely to be fully fixed anytime soon, # so we work around it by preventing an AsyncGenerator and # its _coroutine from being considered newly unreachable at # the same time if the AsyncGenerator's finalizer might want # to iterate the coroutine some more. _pypy_issue2786_workaround = set() def __init__(self, coroutine): self._coroutine = coroutine self._it = coroutine.__await__() self.ag_running = False self._finalizer = None self._closed = False self._hooks_inited = False # On python 3.5.0 and 3.5.1, __aiter__ must be awaitable. # Starting in 3.5.2, it should not be awaitable, and if it is, then it # raises a PendingDeprecationWarning. # See: # https://www.python.org/dev/peps/pep-0492/#api-design-and-implementation-revisions # https://docs.python.org/3/reference/datamodel.html#async-iterators # https://bugs.python.org/issue27243 if sys.version_info < (3, 5, 2): async def __aiter__(self): return self else: def __aiter__(self): return self ################################################################ # Introspection attributes ################################################################ @property def ag_code(self): return self._coroutine.cr_code @property def ag_frame(self): return self._coroutine.cr_frame ################################################################ # Core functionality ################################################################ # These need to return awaitables, rather than being async functions, # to match the native behavior where the firstiter hook is called # immediately on asend()/etc, even if the coroutine that asend() # produces isn't awaited for a bit. def __anext__(self): return self._do_it(self._it.__next__) def asend(self, value): return self._do_it(self._it.send, value) def athrow(self, type, value=None, traceback=None): return self._do_it(self._it.throw, type, value, traceback) def _do_it(self, start_fn, *args): if not self._hooks_inited: self._hooks_inited = True (firstiter, self._finalizer) = get_asyncgen_hooks() if firstiter is not None: firstiter(self) if sys.implementation.name == "pypy": self._pypy_issue2786_workaround.add(self._coroutine) # On CPython 3.5.2 (but not 3.5.0), coroutines get cranky if you try # to iterate them after they're exhausted. Generators OTOH just raise # StopIteration. We want to convert the one into the other, so we need # to avoid iterating stopped coroutines. if getcoroutinestate(self._coroutine) is CORO_CLOSED: raise StopAsyncIteration() async def step(): if self.ag_running: raise ValueError("async generator already executing") try: self.ag_running = True return await ANextIter(self._it, start_fn, *args) except StopAsyncIteration: self._pypy_issue2786_workaround.discard(self._coroutine) raise finally: self.ag_running = False return step() ################################################################ # Cleanup ################################################################ async def aclose(self): state = getcoroutinestate(self._coroutine) if state is CORO_CLOSED or self._closed: return # Make sure that even if we raise "async_generator ignored # GeneratorExit", and thus fail to exhaust the coroutine, # __del__ doesn't complain again. self._closed = True if state is CORO_CREATED: # Make sure that aclose() on an unstarted generator returns # successfully and prevents future iteration. self._it.close() return try: await self.athrow(GeneratorExit) except (GeneratorExit, StopAsyncIteration): self._pypy_issue2786_workaround.discard(self._coroutine) else: raise RuntimeError("async_generator ignored GeneratorExit") def __del__(self): self._pypy_issue2786_workaround.discard(self._coroutine) if getcoroutinestate(self._coroutine) is CORO_CREATED: # Never started, nothing to clean up, just suppress the "coroutine # never awaited" message. self._coroutine.close() if getcoroutinestate(self._coroutine ) is CORO_SUSPENDED and not self._closed: if self._finalizer is not None: self._finalizer(self) else: # Mimic the behavior of native generators on GC with no finalizer: # throw in GeneratorExit, run for one turn, and complain if it didn't # finish. thrower = self.athrow(GeneratorExit) try: thrower.send(None) except (GeneratorExit, StopAsyncIteration): pass except StopIteration: raise RuntimeError("async_generator ignored GeneratorExit") else: raise RuntimeError( "async_generator {!r} awaited during finalization; install " "a finalization hook to support this, or wrap it in " "'async with aclosing(...):'" .format(self.ag_code.co_name) ) finally: thrower.close() if hasattr(collections.abc, "AsyncGenerator"): collections.abc.AsyncGenerator.register(AsyncGenerator) def async_generator(coroutine_maker): @wraps(coroutine_maker) def async_generator_maker(*args, **kwargs): return AsyncGenerator(coroutine_maker(*args, **kwargs)) async_generator_maker._async_gen_function = id(async_generator_maker) return async_generator_maker def isasyncgen(obj): if hasattr(inspect, "isasyncgen"): if inspect.isasyncgen(obj): return True return isinstance(obj, AsyncGenerator) def isasyncgenfunction(obj): if hasattr(inspect, "isasyncgenfunction"): if inspect.isasyncgenfunction(obj): return True return getattr(obj, "_async_gen_function", -1) == id(obj)