456 lines
16 KiB
Python
456 lines
16 KiB
Python
|
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)
|