diff --git a/Include/ceval.h b/Include/ceval.h index 81f4bbf..c682063 100644 --- a/Include/ceval.h +++ b/Include/ceval.h @@ -25,6 +25,10 @@ PyAPI_FUNC(void) PyEval_SetProfile(Py_tracefunc, PyObject *); PyAPI_FUNC(void) PyEval_SetTrace(Py_tracefunc, PyObject *); PyAPI_FUNC(void) _PyEval_SetCoroutineWrapper(PyObject *); PyAPI_FUNC(PyObject *) _PyEval_GetCoroutineWrapper(void); +PyAPI_FUNC(void) _PyEval_SetAsyncGenFirstiter(PyObject *); +PyAPI_FUNC(PyObject *) _PyEval_GetAsyncGenFirstiter(void); +PyAPI_FUNC(void) _PyEval_SetAsyncGenFinalizer(PyObject *); +PyAPI_FUNC(PyObject *) _PyEval_GetAsyncGenFinalizer(void); #endif struct _frame; /* Avoid including frameobject.h */ diff --git a/Include/code.h b/Include/code.h index b39d6bd..9823f10 100644 --- a/Include/code.h +++ b/Include/code.h @@ -59,6 +59,7 @@ typedef struct { ``async def`` keywords) */ #define CO_COROUTINE 0x0080 #define CO_ITERABLE_COROUTINE 0x0100 +#define CO_ASYNC_GENERATOR 0x0200 /* These are no longer used. */ #if 0 diff --git a/Include/genobject.h b/Include/genobject.h index 1ff32a8..973bdd5 100644 --- a/Include/genobject.h +++ b/Include/genobject.h @@ -61,6 +61,37 @@ PyObject *_PyAIterWrapper_New(PyObject *aiter); PyObject *_PyCoro_GetAwaitableIter(PyObject *o); PyAPI_FUNC(PyObject *) PyCoro_New(struct _frame *, PyObject *name, PyObject *qualname); + +/* Asynchronous Generators */ + +typedef struct { + _PyGenObject_HEAD(ag) + PyObject *ag_finalizer; + + /* Flag is set to 1 when hooks set up by sys.set_asyncgen_hooks + were called on the generator, to avoid calling them more + than once. */ + int ag_hooks_inited; + + /* Flag is set to 1 when aclose() is called for the first time, or + when a StopAsyncIteration exception is raised. */ + int ag_closed; +} PyAsyncGenObject; + +PyAPI_DATA(PyTypeObject) PyAsyncGen_Type; +PyAPI_DATA(PyTypeObject) _PyAsyncGenASend_Type; +PyAPI_DATA(PyTypeObject) _PyAsyncGenWrappedValue_Type; +PyAPI_DATA(PyTypeObject) _PyAsyncGenAThrow_Type; + +PyAPI_FUNC(PyObject *) PyAsyncGen_New(struct _frame *, + PyObject *name, PyObject *qualname); + +#define PyAsyncGen_CheckExact(op) (Py_TYPE(op) == &PyAsyncGen_Type) + +PyObject *_PyAsyncGenValueWrapperNew(PyObject *); + +int PyAsyncGen_ClearFreeLists(void); + #endif #undef _PyGenObject_HEAD diff --git a/Include/pylifecycle.h b/Include/pylifecycle.h index cf149b2..8390467 100644 --- a/Include/pylifecycle.h +++ b/Include/pylifecycle.h @@ -107,6 +107,7 @@ PyAPI_FUNC(void) _PyGC_Fini(void); PyAPI_FUNC(void) PySlice_Fini(void); PyAPI_FUNC(void) _PyType_Fini(void); PyAPI_FUNC(void) _PyRandom_Fini(void); +PyAPI_FUNC(void) PyAsyncGen_Fini(void); PyAPI_DATA(PyThreadState *) _Py_Finalizing; #endif diff --git a/Include/pystate.h b/Include/pystate.h index 5ab5c98..f1c9427 100644 --- a/Include/pystate.h +++ b/Include/pystate.h @@ -148,6 +148,9 @@ typedef struct _ts { Py_ssize_t co_extra_user_count; freefunc co_extra_freefuncs[MAX_CO_EXTRA_USERS]; + PyObject *async_gen_firstiter; + PyObject *async_gen_finalizer; + /* XXX signal handlers should also be here */ } PyThreadState; diff --git a/Include/symtable.h b/Include/symtable.h index b0259d6..86ae3c2 100644 --- a/Include/symtable.h +++ b/Include/symtable.h @@ -48,6 +48,7 @@ typedef struct _symtable_entry { unsigned ste_child_free : 1; /* true if a child block has free vars, including free refs to globals */ unsigned ste_generator : 1; /* true if namespace is a generator */ + unsigned ste_coroutine : 1; /* true if namespace is a coroutine */ unsigned ste_varargs : 1; /* true if block has varargs */ unsigned ste_varkeywords : 1; /* true if block has varkeywords */ unsigned ste_returns_value : 1; /* true if namespace uses return with diff --git a/Lib/asyncio/base_events.py b/Lib/asyncio/base_events.py index 918b869..b420586 100644 --- a/Lib/asyncio/base_events.py +++ b/Lib/asyncio/base_events.py @@ -13,7 +13,6 @@ conscious design decision, leaving the door open for keyword arguments to modify the meaning of the API call itself. """ - import collections import concurrent.futures import heapq @@ -28,6 +27,7 @@ import time import traceback import sys import warnings +import weakref from . import compat from . import coroutines @@ -242,6 +242,13 @@ class BaseEventLoop(events.AbstractEventLoop): self._task_factory = None self._coroutine_wrapper_set = False + # A weak set of all asynchronous generators that are being iterated + # by the loop. + self._asyncgens = weakref.WeakSet() + + # Set to True when `loop.shutdown_asyncgens` is called. + self._asyncgens_shutdown_called = False + def __repr__(self): return ('<%s running=%s closed=%s debug=%s>' % (self.__class__.__name__, self.is_running(), @@ -333,6 +340,46 @@ class BaseEventLoop(events.AbstractEventLoop): if self._closed: raise RuntimeError('Event loop is closed') + def _asyncgen_finalizer_hook(self, agen): + self._asyncgens.discard(agen) + if not self.is_closed(): + self.create_task(agen.aclose()) + + def _asyncgen_firstiter_hook(self, agen): + if self._asyncgens_shutdown_called: + warnings.warn( + "asynchronous generator {!r} was scheduled after " + "loop.shutdown_asyncgens() call".format(agen), + ResourceWarning, source=self) + + self._asyncgens.add(agen) + + @coroutine + def shutdown_asyncgens(self): + """Shutdown all active asynchronous generators.""" + self._asyncgens_shutdown_called = True + + if not len(self._asyncgens): + return + + closing_agens = list(self._asyncgens) + self._asyncgens.clear() + + shutdown_coro = tasks.gather( + *[ag.aclose() for ag in closing_agens], + return_exceptions=True, + loop=self) + + results = yield from shutdown_coro + for result, agen in zip(results, closing_agens): + if isinstance(result, Exception): + self.call_exception_handler({ + 'message': 'an error occurred during closing of ' + 'asynchronous generator {!r}'.format(agen), + 'exception': result, + 'asyncgen': agen + }) + def run_forever(self): """Run until stop() is called.""" self._check_closed() @@ -340,6 +387,9 @@ class BaseEventLoop(events.AbstractEventLoop): raise RuntimeError('Event loop is running.') self._set_coroutine_wrapper(self._debug) self._thread_id = threading.get_ident() + old_agen_hooks = sys.get_asyncgen_hooks() + sys.set_asyncgen_hooks(firstiter=self._asyncgen_firstiter_hook, + finalizer=self._asyncgen_finalizer_hook) try: while True: self._run_once() @@ -349,6 +399,7 @@ class BaseEventLoop(events.AbstractEventLoop): self._stopping = False self._thread_id = None self._set_coroutine_wrapper(False) + sys.set_asyncgen_hooks(*old_agen_hooks) def run_until_complete(self, future): """Run until the Future is done. @@ -1179,7 +1230,9 @@ class BaseEventLoop(events.AbstractEventLoop): - 'handle' (optional): Handle instance; - 'protocol' (optional): Protocol instance; - 'transport' (optional): Transport instance; - - 'socket' (optional): Socket instance. + - 'socket' (optional): Socket instance; + - 'asyncgen' (optional): Asynchronous generator that caused + the exception. New keys maybe introduced in the future. diff --git a/Lib/asyncio/coroutines.py b/Lib/asyncio/coroutines.py index 71bc6fb..9c338b0 100644 --- a/Lib/asyncio/coroutines.py +++ b/Lib/asyncio/coroutines.py @@ -276,7 +276,10 @@ def _format_coroutine(coro): try: coro_code = coro.gi_code except AttributeError: - coro_code = coro.cr_code + try: + coro_code = coro.cr_code + except AttributeError: + return repr(coro) try: coro_frame = coro.gi_frame diff --git a/Lib/asyncio/events.py b/Lib/asyncio/events.py index c48c5be..cc9a986 100644 --- a/Lib/asyncio/events.py +++ b/Lib/asyncio/events.py @@ -248,6 +248,10 @@ class AbstractEventLoop: """ raise NotImplementedError + def shutdown_asyncgens(self): + """Shutdown all active asynchronous generators.""" + raise NotImplementedError + # Methods scheduling callbacks. All these return Handles. def _timer_handle_cancelled(self, handle): diff --git a/Lib/dis.py b/Lib/dis.py index 59886f1..556d84e 100644 --- a/Lib/dis.py +++ b/Lib/dis.py @@ -87,6 +87,7 @@ COMPILER_FLAG_NAMES = { 64: "NOFREE", 128: "COROUTINE", 256: "ITERABLE_COROUTINE", + 512: "ASYNC_GENERATOR", } def pretty_flags(flags): diff --git a/Lib/inspect.py b/Lib/inspect.py index 72c1691..2380095 100644 --- a/Lib/inspect.py +++ b/Lib/inspect.py @@ -185,6 +185,13 @@ def iscoroutinefunction(object): return bool((isfunction(object) or ismethod(object)) and object.__code__.co_flags & CO_COROUTINE) +def isasyncgenfunction(object): + return bool((isfunction(object) or ismethod(object)) and + object.__code__.co_flags & CO_ASYNC_GENERATOR) + +def isasyncgen(object): + return isinstance(object, types.AsyncGeneratorType) + def isgenerator(object): """Return true if the object is a generator. diff --git a/Lib/test/badsyntax_async6.py b/Lib/test/badsyntax_async6.py deleted file mode 100644 index cb0a23d..0000000 --- a/Lib/test/badsyntax_async6.py +++ /dev/null @@ -1,2 +0,0 @@ -async def foo(): - yield diff --git a/Lib/test/test_asyncgen.py b/Lib/test/test_asyncgen.py new file mode 100644 index 0000000..41b1b4f --- /dev/null +++ b/Lib/test/test_asyncgen.py @@ -0,0 +1,823 @@ +import asyncio +import inspect +import sys +import types +import unittest + +from unittest import mock + + +class AwaitException(Exception): + pass + + +@types.coroutine +def awaitable(*, throw=False): + if throw: + yield ('throw',) + else: + yield ('result',) + + +def run_until_complete(coro): + exc = False + while True: + try: + if exc: + exc = False + fut = coro.throw(AwaitException) + else: + fut = coro.send(None) + except StopIteration as ex: + return ex.args[0] + + if fut == ('throw',): + exc = True + + +def to_list(gen): + async def iterate(): + res = [] + async for i in gen: + res.append(i) + return res + + return run_until_complete(iterate()) + + +class AsyncGenSyntaxTest(unittest.TestCase): + + def test_async_gen_syntax_01(self): + code = '''async def foo(): + await abc + yield from 123 + ''' + + with self.assertRaisesRegex(SyntaxError, 'yield from.*inside async'): + exec(code, {}, {}) + + def test_async_gen_syntax_02(self): + code = '''async def foo(): + yield from 123 + ''' + + with self.assertRaisesRegex(SyntaxError, 'yield from.*inside async'): + exec(code, {}, {}) + + def test_async_gen_syntax_03(self): + code = '''async def foo(): + await abc + yield + return 123 + ''' + + with self.assertRaisesRegex(SyntaxError, 'return.*value.*async gen'): + exec(code, {}, {}) + + def test_async_gen_syntax_04(self): + code = '''async def foo(): + yield + return 123 + ''' + + with self.assertRaisesRegex(SyntaxError, 'return.*value.*async gen'): + exec(code, {}, {}) + + def test_async_gen_syntax_05(self): + code = '''async def foo(): + if 0: + yield + return 12 + ''' + + with self.assertRaisesRegex(SyntaxError, 'return.*value.*async gen'): + exec(code, {}, {}) + + +class AsyncGenTest(unittest.TestCase): + + def compare_generators(self, sync_gen, async_gen): + def sync_iterate(g): + res = [] + while True: + try: + res.append(g.__next__()) + except StopIteration: + res.append('STOP') + break + except Exception as ex: + res.append(str(type(ex))) + return res + + def async_iterate(g): + res = [] + while True: + try: + g.__anext__().__next__() + except StopAsyncIteration: + res.append('STOP') + break + except StopIteration as ex: + if ex.args: + res.append(ex.args[0]) + else: + res.append('EMPTY StopIteration') + break + except Exception as ex: + res.append(str(type(ex))) + return res + + sync_gen_result = sync_iterate(sync_gen) + async_gen_result = async_iterate(async_gen) + self.assertEqual(sync_gen_result, async_gen_result) + return async_gen_result + + def test_async_gen_iteration_01(self): + async def gen(): + await awaitable() + a = yield 123 + self.assertIs(a, None) + await awaitable() + yield 456 + await awaitable() + yield 789 + + self.assertEqual(to_list(gen()), [123, 456, 789]) + + def test_async_gen_iteration_02(self): + async def gen(): + await awaitable() + yield 123 + await awaitable() + + g = gen() + ai = g.__aiter__() + self.assertEqual(ai.__anext__().__next__(), ('result',)) + + try: + ai.__anext__().__next__() + except StopIteration as ex: + self.assertEqual(ex.args[0], 123) + else: + self.fail('StopIteration was not raised') + + self.assertEqual(ai.__anext__().__next__(), ('result',)) + + try: + ai.__anext__().__next__() + except StopAsyncIteration as ex: + self.assertFalse(ex.args) + else: + self.fail('StopAsyncIteration was not raised') + + def test_async_gen_exception_03(self): + async def gen(): + await awaitable() + yield 123 + await awaitable(throw=True) + yield 456 + + with self.assertRaises(AwaitException): + to_list(gen()) + + def test_async_gen_exception_04(self): + async def gen(): + await awaitable() + yield 123 + 1 / 0 + + g = gen() + ai = g.__aiter__() + self.assertEqual(ai.__anext__().__next__(), ('result',)) + + try: + ai.__anext__().__next__() + except StopIteration as ex: + self.assertEqual(ex.args[0], 123) + else: + self.fail('StopIteration was not raised') + + with self.assertRaises(ZeroDivisionError): + ai.__anext__().__next__() + + def test_async_gen_exception_05(self): + async def gen(): + yield 123 + raise StopAsyncIteration + + with self.assertRaisesRegex(RuntimeError, + 'async generator.*StopAsyncIteration'): + to_list(gen()) + + def test_async_gen_exception_06(self): + async def gen(): + yield 123 + raise StopIteration + + with self.assertRaisesRegex(RuntimeError, + 'async generator.*StopIteration'): + to_list(gen()) + + def test_async_gen_exception_07(self): + def sync_gen(): + try: + yield 1 + 1 / 0 + finally: + yield 2 + yield 3 + + yield 100 + + async def async_gen(): + try: + yield 1 + 1 / 0 + finally: + yield 2 + yield 3 + + yield 100 + + self.compare_generators(sync_gen(), async_gen()) + + def test_async_gen_exception_08(self): + def sync_gen(): + try: + yield 1 + finally: + yield 2 + 1 / 0 + yield 3 + + yield 100 + + async def async_gen(): + try: + yield 1 + await awaitable() + finally: + await awaitable() + yield 2 + 1 / 0 + yield 3 + + yield 100 + + self.compare_generators(sync_gen(), async_gen()) + + def test_async_gen_exception_09(self): + def sync_gen(): + try: + yield 1 + 1 / 0 + finally: + yield 2 + yield 3 + + yield 100 + + async def async_gen(): + try: + await awaitable() + yield 1 + 1 / 0 + finally: + yield 2 + await awaitable() + yield 3 + + yield 100 + + self.compare_generators(sync_gen(), async_gen()) + + def test_async_gen_exception_10(self): + async def gen(): + yield 123 + with self.assertRaisesRegex(TypeError, + "non-None value .* async generator"): + gen().__anext__().send(100) + + def test_async_gen_api_01(self): + async def gen(): + yield 123 + + g = gen() + + self.assertEqual(g.__name__, 'gen') + g.__name__ = '123' + self.assertEqual(g.__name__, '123') + + self.assertIn('.gen', g.__qualname__) + g.__qualname__ = '123' + self.assertEqual(g.__qualname__, '123') + + self.assertIsNone(g.ag_await) + self.assertIsInstance(g.ag_frame, types.FrameType) + self.assertFalse(g.ag_running) + self.assertIsInstance(g.ag_code, types.CodeType) + + self.assertTrue(inspect.isawaitable(g.aclose())) + + +class AsyncGenAsyncioTest(unittest.TestCase): + + def setUp(self): + self.loop = asyncio.new_event_loop() + asyncio.set_event_loop(None) + + def tearDown(self): + self.loop.close() + self.loop = None + + async def to_list(self, gen): + res = [] + async for i in gen: + res.append(i) + return res + + def test_async_gen_asyncio_01(self): + async def gen(): + yield 1 + await asyncio.sleep(0.01, loop=self.loop) + yield 2 + await asyncio.sleep(0.01, loop=self.loop) + return + yield 3 + + res = self.loop.run_until_complete(self.to_list(gen())) + self.assertEqual(res, [1, 2]) + + def test_async_gen_asyncio_02(self): + async def gen(): + yield 1 + await asyncio.sleep(0.01, loop=self.loop) + yield 2 + 1 / 0 + yield 3 + + with self.assertRaises(ZeroDivisionError): + self.loop.run_until_complete(self.to_list(gen())) + + def test_async_gen_asyncio_03(self): + loop = self.loop + + class Gen: + async def __aiter__(self): + yield 1 + await asyncio.sleep(0.01, loop=loop) + yield 2 + + res = loop.run_until_complete(self.to_list(Gen())) + self.assertEqual(res, [1, 2]) + + def test_async_gen_asyncio_anext_04(self): + async def foo(): + yield 1 + await asyncio.sleep(0.01, loop=self.loop) + try: + yield 2 + yield 3 + except ZeroDivisionError: + yield 1000 + await asyncio.sleep(0.01, loop=self.loop) + yield 4 + + async def run1(): + it = foo().__aiter__() + + self.assertEqual(await it.__anext__(), 1) + self.assertEqual(await it.__anext__(), 2) + self.assertEqual(await it.__anext__(), 3) + self.assertEqual(await it.__anext__(), 4) + with self.assertRaises(StopAsyncIteration): + await it.__anext__() + with self.assertRaises(StopAsyncIteration): + await it.__anext__() + + async def run2(): + it = foo().__aiter__() + + self.assertEqual(await it.__anext__(), 1) + self.assertEqual(await it.__anext__(), 2) + try: + it.__anext__().throw(ZeroDivisionError) + except StopIteration as ex: + self.assertEqual(ex.args[0], 1000) + else: + self.fail('StopIteration was not raised') + self.assertEqual(await it.__anext__(), 4) + with self.assertRaises(StopAsyncIteration): + await it.__anext__() + + self.loop.run_until_complete(run1()) + self.loop.run_until_complete(run2()) + + def test_async_gen_asyncio_anext_05(self): + async def foo(): + v = yield 1 + v = yield v + yield v * 100 + + async def run(): + it = foo().__aiter__() + + try: + it.__anext__().send(None) + except StopIteration as ex: + self.assertEqual(ex.args[0], 1) + else: + self.fail('StopIteration was not raised') + + try: + it.__anext__().send(10) + except StopIteration as ex: + self.assertEqual(ex.args[0], 10) + else: + self.fail('StopIteration was not raised') + + try: + it.__anext__().send(12) + except StopIteration as ex: + self.assertEqual(ex.args[0], 1200) + else: + self.fail('StopIteration was not raised') + + with self.assertRaises(StopAsyncIteration): + await it.__anext__() + + self.loop.run_until_complete(run()) + + def test_async_gen_asyncio_aclose_06(self): + async def foo(): + try: + yield 1 + 1 / 0 + finally: + await asyncio.sleep(0.01, loop=self.loop) + yield 12 + + async def run(): + gen = foo() + it = gen.__aiter__() + await it.__anext__() + await gen.aclose() + + with self.assertRaisesRegex( + RuntimeError, + "async generator ignored GeneratorExit"): + self.loop.run_until_complete(run()) + + def test_async_gen_asyncio_aclose_07(self): + DONE = 0 + + async def foo(): + nonlocal DONE + try: + yield 1 + 1 / 0 + finally: + await asyncio.sleep(0.01, loop=self.loop) + await asyncio.sleep(0.01, loop=self.loop) + DONE += 1 + DONE += 1000 + + async def run(): + gen = foo() + it = gen.__aiter__() + await it.__anext__() + await gen.aclose() + + self.loop.run_until_complete(run()) + self.assertEqual(DONE, 1) + + def test_async_gen_asyncio_aclose_08(self): + DONE = 0 + + fut = asyncio.Future(loop=self.loop) + + async def foo(): + nonlocal DONE + try: + yield 1 + await fut + DONE += 1000 + yield 2 + finally: + await asyncio.sleep(0.01, loop=self.loop) + await asyncio.sleep(0.01, loop=self.loop) + DONE += 1 + DONE += 1000 + + async def run(): + gen = foo() + it = gen.__aiter__() + self.assertEqual(await it.__anext__(), 1) + t = self.loop.create_task(it.__anext__()) + await asyncio.sleep(0.01, loop=self.loop) + await gen.aclose() + return t + + t = self.loop.run_until_complete(run()) + self.assertEqual(DONE, 1) + + # Silence ResourceWarnings + fut.cancel() + t.cancel() + self.loop.run_until_complete(asyncio.sleep(0.01, loop=self.loop)) + + def test_async_gen_asyncio_gc_aclose_09(self): + DONE = 0 + + async def gen(): + nonlocal DONE + try: + while True: + yield 1 + finally: + await asyncio.sleep(0.01, loop=self.loop) + await asyncio.sleep(0.01, loop=self.loop) + DONE = 1 + + async def run(): + g = gen() + await g.__anext__() + await g.__anext__() + del g + + await asyncio.sleep(0.1, loop=self.loop) + + self.loop.run_until_complete(run()) + self.assertEqual(DONE, 1) + + def test_async_gen_asyncio_asend_01(self): + DONE = 0 + + # Sanity check: + def sgen(): + v = yield 1 + yield v * 2 + sg = sgen() + v = sg.send(None) + self.assertEqual(v, 1) + v = sg.send(100) + self.assertEqual(v, 200) + + async def gen(): + nonlocal DONE + try: + await asyncio.sleep(0.01, loop=self.loop) + v = yield 1 + await asyncio.sleep(0.01, loop=self.loop) + yield v * 2 + await asyncio.sleep(0.01, loop=self.loop) + return + finally: + await asyncio.sleep(0.01, loop=self.loop) + await asyncio.sleep(0.01, loop=self.loop) + DONE = 1 + + async def run(): + g = gen() + + v = await g.asend(None) + self.assertEqual(v, 1) + + v = await g.asend(100) + self.assertEqual(v, 200) + + with self.assertRaises(StopAsyncIteration): + await g.asend(None) + + self.loop.run_until_complete(run()) + self.assertEqual(DONE, 1) + + def test_async_gen_asyncio_asend_02(self): + DONE = 0 + + async def sleep_n_crash(delay): + await asyncio.sleep(delay, loop=self.loop) + 1 / 0 + + async def gen(): + nonlocal DONE + try: + await asyncio.sleep(0.01, loop=self.loop) + v = yield 1 + await sleep_n_crash(0.01) + DONE += 1000 + yield v * 2 + finally: + await asyncio.sleep(0.01, loop=self.loop) + await asyncio.sleep(0.01, loop=self.loop) + DONE = 1 + + async def run(): + g = gen() + + v = await g.asend(None) + self.assertEqual(v, 1) + + await g.asend(100) + + with self.assertRaises(ZeroDivisionError): + self.loop.run_until_complete(run()) + self.assertEqual(DONE, 1) + + def test_async_gen_asyncio_asend_03(self): + DONE = 0 + + async def sleep_n_crash(delay): + fut = asyncio.ensure_future(asyncio.sleep(delay, loop=self.loop), + loop=self.loop) + self.loop.call_later(delay / 2, lambda: fut.cancel()) + return await fut + + async def gen(): + nonlocal DONE + try: + await asyncio.sleep(0.01, loop=self.loop) + v = yield 1 + await sleep_n_crash(0.01) + DONE += 1000 + yield v * 2 + finally: + await asyncio.sleep(0.01, loop=self.loop) + await asyncio.sleep(0.01, loop=self.loop) + DONE = 1 + + async def run(): + g = gen() + + v = await g.asend(None) + self.assertEqual(v, 1) + + await g.asend(100) + + with self.assertRaises(asyncio.CancelledError): + self.loop.run_until_complete(run()) + self.assertEqual(DONE, 1) + + def test_async_gen_asyncio_athrow_01(self): + DONE = 0 + + class FooEr(Exception): + pass + + # Sanity check: + def sgen(): + try: + v = yield 1 + except FooEr: + v = 1000 + yield v * 2 + sg = sgen() + v = sg.send(None) + self.assertEqual(v, 1) + v = sg.throw(FooEr) + self.assertEqual(v, 2000) + with self.assertRaises(StopIteration): + sg.send(None) + + async def gen(): + nonlocal DONE + try: + await asyncio.sleep(0.01, loop=self.loop) + try: + v = yield 1 + except FooEr: + v = 1000 + await asyncio.sleep(0.01, loop=self.loop) + yield v * 2 + await asyncio.sleep(0.01, loop=self.loop) + # return + finally: + await asyncio.sleep(0.01, loop=self.loop) + await asyncio.sleep(0.01, loop=self.loop) + DONE = 1 + + async def run(): + g = gen() + + v = await g.asend(None) + self.assertEqual(v, 1) + + v = await g.athrow(FooEr) + self.assertEqual(v, 2000) + + with self.assertRaises(StopAsyncIteration): + await g.asend(None) + + self.loop.run_until_complete(run()) + self.assertEqual(DONE, 1) + + def test_async_gen_asyncio_athrow_02(self): + DONE = 0 + + class FooEr(Exception): + pass + + async def sleep_n_crash(delay): + fut = asyncio.ensure_future(asyncio.sleep(delay, loop=self.loop), + loop=self.loop) + self.loop.call_later(delay / 2, lambda: fut.cancel()) + return await fut + + async def gen(): + nonlocal DONE + try: + await asyncio.sleep(0.01, loop=self.loop) + try: + v = yield 1 + except FooEr: + await sleep_n_crash(0.01) + yield v * 2 + await asyncio.sleep(0.01, loop=self.loop) + # return + finally: + await asyncio.sleep(0.01, loop=self.loop) + await asyncio.sleep(0.01, loop=self.loop) + DONE = 1 + + async def run(): + g = gen() + + v = await g.asend(None) + self.assertEqual(v, 1) + + try: + await g.athrow(FooEr) + except asyncio.CancelledError: + self.assertEqual(DONE, 1) + raise + else: + self.fail('CancelledError was not raised') + + with self.assertRaises(asyncio.CancelledError): + self.loop.run_until_complete(run()) + self.assertEqual(DONE, 1) + + def test_async_gen_asyncio_shutdown_01(self): + finalized = 0 + + async def waiter(timeout): + nonlocal finalized + try: + await asyncio.sleep(timeout, loop=self.loop) + yield 1 + finally: + await asyncio.sleep(0, loop=self.loop) + finalized += 1 + + async def wait(): + async for _ in waiter(1): + pass + + t1 = self.loop.create_task(wait()) + t2 = self.loop.create_task(wait()) + + self.loop.run_until_complete(asyncio.sleep(0.1, loop=self.loop)) + + self.loop.run_until_complete(self.loop.shutdown_asyncgens()) + self.assertEqual(finalized, 2) + + # Silence warnings + t1.cancel() + t2.cancel() + self.loop.run_until_complete(asyncio.sleep(0.1, loop=self.loop)) + + def test_async_gen_asyncio_shutdown_02(self): + logged = 0 + + def logger(loop, context): + nonlocal logged + self.assertIn('asyncgen', context) + expected = 'an error occurred during closing of asynchronous' + if expected in context['message']: + logged += 1 + + async def waiter(timeout): + try: + await asyncio.sleep(timeout, loop=self.loop) + yield 1 + finally: + 1 / 0 + + async def wait(): + async for _ in waiter(1): + pass + + t = self.loop.create_task(wait()) + self.loop.run_until_complete(asyncio.sleep(0.1, loop=self.loop)) + + self.loop.set_exception_handler(logger) + self.loop.run_until_complete(self.loop.shutdown_asyncgens()) + + self.assertEqual(logged, 1) + + # Silence warnings + t.cancel() + self.loop.run_until_complete(asyncio.sleep(0.1, loop=self.loop)) + +if __name__ == "__main__": + unittest.main() diff --git a/Lib/test/test_coroutines.py b/Lib/test/test_coroutines.py index ee2482d..fee9ae3 100644 --- a/Lib/test/test_coroutines.py +++ b/Lib/test/test_coroutines.py @@ -88,12 +88,6 @@ class AsyncBadSyntaxTest(unittest.TestCase): with self.assertRaisesRegex(SyntaxError, 'invalid syntax'): import test.badsyntax_async5 - def test_badsyntax_6(self): - with self.assertRaisesRegex( - SyntaxError, "'yield' inside async function"): - - import test.badsyntax_async6 - def test_badsyntax_7(self): with self.assertRaisesRegex( SyntaxError, "'yield from' inside async function"): diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py index 6081073..21b8cb7 100644 --- a/Lib/test/test_dis.py +++ b/Lib/test/test_dis.py @@ -574,7 +574,7 @@ Argument count: 0 Kw-only arguments: 0 Number of locals: 2 Stack size: 17 -Flags: OPTIMIZED, NEWLOCALS, GENERATOR, NOFREE, COROUTINE +Flags: OPTIMIZED, NEWLOCALS, NOFREE, COROUTINE Constants: 0: None 1: 1""" diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py index 47244ae..97634e5 100644 --- a/Lib/test/test_inspect.py +++ b/Lib/test/test_inspect.py @@ -65,7 +65,8 @@ class IsTestBase(unittest.TestCase): inspect.isframe, inspect.isfunction, inspect.ismethod, inspect.ismodule, inspect.istraceback, inspect.isgenerator, inspect.isgeneratorfunction, - inspect.iscoroutine, inspect.iscoroutinefunction]) + inspect.iscoroutine, inspect.iscoroutinefunction, + inspect.isasyncgen, inspect.isasyncgenfunction]) def istest(self, predicate, exp): obj = eval(exp) @@ -73,6 +74,7 @@ class IsTestBase(unittest.TestCase): for other in self.predicates - set([predicate]): if (predicate == inspect.isgeneratorfunction or \ + predicate == inspect.isasyncgenfunction or \ predicate == inspect.iscoroutinefunction) and \ other == inspect.isfunction: continue @@ -82,6 +84,10 @@ def generator_function_example(self): for i in range(2): yield i +async def async_generator_function_example(self): + async for i in range(2): + yield i + async def coroutine_function_example(self): return 'spam' @@ -122,6 +128,10 @@ class TestPredicates(IsTestBase): self.istest(inspect.isdatadescriptor, 'collections.defaultdict.default_factory') self.istest(inspect.isgenerator, '(x for x in range(2))') self.istest(inspect.isgeneratorfunction, 'generator_function_example') + self.istest(inspect.isasyncgen, + 'async_generator_function_example(1)') + self.istest(inspect.isasyncgenfunction, + 'async_generator_function_example') with warnings.catch_warnings(): warnings.simplefilter("ignore") diff --git a/Lib/test/test_sys.py b/Lib/test/test_sys.py index a6cd95e..37ee0b0 100644 --- a/Lib/test/test_sys.py +++ b/Lib/test/test_sys.py @@ -1192,6 +1192,32 @@ class SizeofTest(unittest.TestCase): # sys.flags check(sys.flags, vsize('') + self.P * len(sys.flags)) + def test_asyncgen_hooks(self): + old = sys.get_asyncgen_hooks() + self.assertIsNone(old.firstiter) + self.assertIsNone(old.finalizer) + + firstiter = lambda *a: None + sys.set_asyncgen_hooks(firstiter=firstiter) + hooks = sys.get_asyncgen_hooks() + self.assertIs(hooks.firstiter, firstiter) + self.assertIs(hooks[0], firstiter) + self.assertIs(hooks.finalizer, None) + self.assertIs(hooks[1], None) + + finalizer = lambda *a: None + sys.set_asyncgen_hooks(finalizer=finalizer) + hooks = sys.get_asyncgen_hooks() + self.assertIs(hooks.firstiter, firstiter) + self.assertIs(hooks[0], firstiter) + self.assertIs(hooks.finalizer, finalizer) + self.assertIs(hooks[1], finalizer) + + sys.set_asyncgen_hooks(*old) + cur = sys.get_asyncgen_hooks() + self.assertIsNone(cur.firstiter) + self.assertIsNone(cur.finalizer) + def test_main(): test.support.run_unittest(SysModuleTest, SizeofTest) diff --git a/Lib/types.py b/Lib/types.py index 48891cd..d8d8470 100644 --- a/Lib/types.py +++ b/Lib/types.py @@ -24,6 +24,11 @@ _c = _c() CoroutineType = type(_c) _c.close() # Prevent ResourceWarning +async def _ag(): + yield +_ag = _ag() +AsyncGeneratorType = type(_ag) + class _C: def _m(self): pass MethodType = type(_C()._m) diff --git a/Modules/gcmodule.c b/Modules/gcmodule.c index 0c6f444..07950a6 100644 --- a/Modules/gcmodule.c +++ b/Modules/gcmodule.c @@ -892,6 +892,7 @@ clear_freelists(void) (void)PyList_ClearFreeList(); (void)PyDict_ClearFreeList(); (void)PySet_ClearFreeList(); + (void)PyAsyncGen_ClearFreeLists(); } /* This is the main function. Read this to understand how the diff --git a/Objects/genobject.c b/Objects/genobject.c index 19d388e..bc5309a 100644 --- a/Objects/genobject.c +++ b/Objects/genobject.c @@ -5,7 +5,15 @@ #include "structmember.h" #include "opcode.h" -static PyObject *gen_close(PyGenObject *gen, PyObject *args); +static PyObject *gen_close(PyGenObject *, PyObject *); +static PyObject *async_gen_asend_new(PyAsyncGenObject *, PyObject *); +static PyObject *async_gen_athrow_new(PyAsyncGenObject *, PyObject *); + +static char *NON_INIT_CORO_MSG = "can't send non-None value to a " + "just-started coroutine"; + +static char *ASYNC_GEN_IGNORED_EXIT_MSG = + "async generator ignored GeneratorExit"; static int gen_traverse(PyGenObject *gen, visitproc visit, void *arg) @@ -28,6 +36,26 @@ _PyGen_Finalize(PyObject *self) /* Generator isn't paused, so no need to close */ return; + if (PyAsyncGen_CheckExact(self)) { + PyAsyncGenObject *agen = (PyAsyncGenObject*)self; + PyObject *finalizer = agen->ag_finalizer; + if (finalizer && !agen->ag_closed) { + /* Save the current exception, if any. */ + PyErr_Fetch(&error_type, &error_value, &error_traceback); + + res = PyObject_CallFunctionObjArgs(finalizer, self, NULL); + + if (res == NULL) { + PyErr_WriteUnraisable(self); + } else { + Py_DECREF(res); + } + /* Restore the saved exception. */ + PyErr_Restore(error_type, error_value, error_traceback); + return; + } + } + /* Save the current exception, if any. */ PyErr_Fetch(&error_type, &error_value, &error_traceback); @@ -74,6 +102,12 @@ gen_dealloc(PyGenObject *gen) return; /* resurrected. :( */ _PyObject_GC_UNTRACK(self); + if (PyAsyncGen_CheckExact(gen)) { + /* We have to handle this case for asynchronous generators + right here, because this code has to be between UNTRACK + and GC_Del. */ + Py_CLEAR(((PyAsyncGenObject*)gen)->ag_finalizer); + } if (gen->gi_frame != NULL) { gen->gi_frame->f_gen = NULL; Py_CLEAR(gen->gi_frame); @@ -84,6 +118,33 @@ gen_dealloc(PyGenObject *gen) PyObject_GC_Del(gen); } +static void +gen_chain_runtime_error(const char *msg) +{ + PyObject *exc, *val, *val2, *tb; + + /* TODO: This about rewriting using _PyErr_ChainExceptions. */ + + PyErr_Fetch(&exc, &val, &tb); + PyErr_NormalizeException(&exc, &val, &tb); + if (tb != NULL) { + PyException_SetTraceback(val, tb); + } + + Py_DECREF(exc); + Py_XDECREF(tb); + + PyErr_SetString(PyExc_RuntimeError, msg); + PyErr_Fetch(&exc, &val2, &tb); + PyErr_NormalizeException(&exc, &val2, &tb); + + Py_INCREF(val); + PyException_SetCause(val2, val); + PyException_SetContext(val2, val); + + PyErr_Restore(exc, val2, tb); +} + static PyObject * gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing) { @@ -93,8 +154,12 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing) if (gen->gi_running) { char *msg = "generator already executing"; - if (PyCoro_CheckExact(gen)) + if (PyCoro_CheckExact(gen)) { msg = "coroutine already executing"; + } + else if (PyAsyncGen_CheckExact(gen)) { + msg = "async generator already executing"; + } PyErr_SetString(PyExc_ValueError, msg); return NULL; } @@ -106,10 +171,16 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing) PyErr_SetString( PyExc_RuntimeError, "cannot reuse already awaited coroutine"); - } else if (arg && !exc) { + } + else if (arg && !exc) { /* `gen` is an exhausted generator: only set exception if called from send(). */ - PyErr_SetNone(PyExc_StopIteration); + if (PyAsyncGen_CheckExact(gen)) { + PyErr_SetNone(PyExc_StopAsyncIteration); + } + else { + PyErr_SetNone(PyExc_StopIteration); + } } return NULL; } @@ -118,9 +189,13 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing) if (arg && arg != Py_None) { char *msg = "can't send non-None value to a " "just-started generator"; - if (PyCoro_CheckExact(gen)) + if (PyCoro_CheckExact(gen)) { + msg = NON_INIT_CORO_MSG; + } + else if (PyAsyncGen_CheckExact(gen)) { msg = "can't send non-None value to a " - "just-started coroutine"; + "just-started async generator"; + } PyErr_SetString(PyExc_TypeError, msg); return NULL; } @@ -152,10 +227,20 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing) if (result && f->f_stacktop == NULL) { if (result == Py_None) { /* Delay exception instantiation if we can */ - PyErr_SetNone(PyExc_StopIteration); - } else { + if (PyAsyncGen_CheckExact(gen)) { + PyErr_SetNone(PyExc_StopAsyncIteration); + } + else { + PyErr_SetNone(PyExc_StopIteration); + } + } + else { PyObject *e = PyObject_CallFunctionObjArgs( PyExc_StopIteration, result, NULL); + + /* Async generators cannot return anything but None */ + assert(!PyAsyncGen_CheckExact(gen)); + if (e != NULL) { PyErr_SetObject(PyExc_StopIteration, e); Py_DECREF(e); @@ -167,28 +252,38 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing) /* Check for __future__ generator_stop and conditionally turn * a leaking StopIteration into RuntimeError (with its cause * set appropriately). */ - if (gen->gi_code != NULL && ((PyCodeObject *)gen->gi_code)->co_flags & - (CO_FUTURE_GENERATOR_STOP | CO_COROUTINE | CO_ITERABLE_COROUTINE)) + + const int check_stop_iter_error_flags = CO_FUTURE_GENERATOR_STOP | + CO_COROUTINE | + CO_ITERABLE_COROUTINE | + CO_ASYNC_GENERATOR; + + if (gen->gi_code != NULL && + ((PyCodeObject *)gen->gi_code)->co_flags & + check_stop_iter_error_flags) { - PyObject *exc, *val, *val2, *tb; - char *msg = "generator raised StopIteration"; - if (PyCoro_CheckExact(gen)) + /* `gen` is either: + * a generator with CO_FUTURE_GENERATOR_STOP flag; + * a coroutine; + * a generator with CO_ITERABLE_COROUTINE flag + (decorated with types.coroutine decorator); + * an async generator. + */ + const char *msg = "generator raised StopIteration"; + if (PyCoro_CheckExact(gen)) { msg = "coroutine raised StopIteration"; - PyErr_Fetch(&exc, &val, &tb); - PyErr_NormalizeException(&exc, &val, &tb); - if (tb != NULL) - PyException_SetTraceback(val, tb); - Py_DECREF(exc); - Py_XDECREF(tb); - PyErr_SetString(PyExc_RuntimeError, msg); - PyErr_Fetch(&exc, &val2, &tb); - PyErr_NormalizeException(&exc, &val2, &tb); - Py_INCREF(val); - PyException_SetCause(val2, val); - PyException_SetContext(val2, val); - PyErr_Restore(exc, val2, tb); + } + else if PyAsyncGen_CheckExact(gen) { + msg = "async generator raised StopIteration"; + } + /* Raise a RuntimeError */ + gen_chain_runtime_error(msg); } else { + /* `gen` is an ordinary generator without + CO_FUTURE_GENERATOR_STOP flag. + */ + PyObject *exc, *val, *tb; /* Pop the exception before issuing a warning. */ @@ -207,6 +302,15 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing) } } } + else if (PyAsyncGen_CheckExact(gen) && !result && + PyErr_ExceptionMatches(PyExc_StopAsyncIteration)) + { + /* code in `gen` raised a StopAsyncIteration error: + raise a RuntimeError. + */ + const char *msg = "async generator raised StopAsyncIteration"; + gen_chain_runtime_error(msg); + } if (!result || f->f_stacktop == NULL) { /* generator can't be rerun, so release the frame */ @@ -253,17 +357,19 @@ gen_close_iter(PyObject *yf) PyObject *retval = NULL; _Py_IDENTIFIER(close); - if (PyGen_CheckExact(yf)) { + if (PyGen_CheckExact(yf) || PyCoro_CheckExact(yf)) { retval = gen_close((PyGenObject *)yf, NULL); if (retval == NULL) return -1; - } else { + } + else { PyObject *meth = _PyObject_GetAttrId(yf, &PyId_close); if (meth == NULL) { if (!PyErr_ExceptionMatches(PyExc_AttributeError)) PyErr_WriteUnraisable(yf); PyErr_Clear(); - } else { + } + else { retval = _PyObject_CallNoArg(meth); Py_DECREF(meth); if (retval == NULL) @@ -311,8 +417,11 @@ gen_close(PyGenObject *gen, PyObject *args) retval = gen_send_ex(gen, Py_None, 1, 1); if (retval) { char *msg = "generator ignored GeneratorExit"; - if (PyCoro_CheckExact(gen)) + if (PyCoro_CheckExact(gen)) { msg = "coroutine ignored GeneratorExit"; + } else if (PyAsyncGen_CheckExact(gen)) { + msg = ASYNC_GEN_IGNORED_EXIT_MSG; + } Py_DECREF(retval); PyErr_SetString(PyExc_RuntimeError, msg); return NULL; @@ -332,21 +441,22 @@ PyDoc_STRVAR(throw_doc, return next yielded value or raise StopIteration."); static PyObject * -gen_throw(PyGenObject *gen, PyObject *args) +_gen_throw(PyGenObject *gen, int close_on_genexit, + PyObject *typ, PyObject *val, PyObject *tb) { - PyObject *typ; - PyObject *tb = NULL; - PyObject *val = NULL; PyObject *yf = _PyGen_yf(gen); _Py_IDENTIFIER(throw); - if (!PyArg_UnpackTuple(args, "throw", 1, 3, &typ, &val, &tb)) - return NULL; - if (yf) { PyObject *ret; int err; - if (PyErr_GivenExceptionMatches(typ, PyExc_GeneratorExit)) { + if (PyErr_GivenExceptionMatches(typ, PyExc_GeneratorExit) && + close_on_genexit + ) { + /* Asynchronous generators *should not* be closed right away. + We have to allow some awaits to work it through, hence the + `close_on_genexit` parameter here. + */ gen->gi_running = 1; err = gen_close_iter(yf); gen->gi_running = 0; @@ -355,11 +465,16 @@ gen_throw(PyGenObject *gen, PyObject *args) return gen_send_ex(gen, Py_None, 1, 0); goto throw_here; } - if (PyGen_CheckExact(yf)) { + if (PyGen_CheckExact(yf) || PyCoro_CheckExact(yf)) { + /* `yf` is a generator or a coroutine. */ gen->gi_running = 1; - ret = gen_throw((PyGenObject *)yf, args); + /* Close the generator that we are currently iterating with + 'yield from' or awaiting on with 'await'. */ + ret = _gen_throw((PyGenObject *)yf, close_on_genexit, + typ, val, tb); gen->gi_running = 0; } else { + /* `yf` is an iterator or a coroutine-like object. */ PyObject *meth = _PyObject_GetAttrId(yf, &PyId_throw); if (meth == NULL) { if (!PyErr_ExceptionMatches(PyExc_AttributeError)) { @@ -371,7 +486,7 @@ gen_throw(PyGenObject *gen, PyObject *args) goto throw_here; } gen->gi_running = 1; - ret = PyObject_CallObject(meth, args); + ret = PyObject_CallFunctionObjArgs(meth, typ, val, tb, NULL); gen->gi_running = 0; Py_DECREF(meth); } @@ -454,6 +569,21 @@ failed_throw: static PyObject * +gen_throw(PyGenObject *gen, PyObject *args) +{ + PyObject *typ; + PyObject *tb = NULL; + PyObject *val = NULL; + + if (!PyArg_UnpackTuple(args, "throw", 1, 3, &typ, &val, &tb)) { + return NULL; + } + + return _gen_throw(gen, 1, typ, val, tb); +} + + +static PyObject * gen_iternext(PyGenObject *gen) { return gen_send_ex(gen, NULL, 0, 0); @@ -997,21 +1127,21 @@ PyCoro_New(PyFrameObject *f, PyObject *name, PyObject *qualname) typedef struct { PyObject_HEAD - PyObject *aw_aiter; + PyObject *ags_aiter; } PyAIterWrapper; static PyObject * aiter_wrapper_iternext(PyAIterWrapper *aw) { - PyErr_SetObject(PyExc_StopIteration, aw->aw_aiter); + PyErr_SetObject(PyExc_StopIteration, aw->ags_aiter); return NULL; } static int aiter_wrapper_traverse(PyAIterWrapper *aw, visitproc visit, void *arg) { - Py_VISIT((PyObject *)aw->aw_aiter); + Py_VISIT((PyObject *)aw->ags_aiter); return 0; } @@ -1019,7 +1149,7 @@ static void aiter_wrapper_dealloc(PyAIterWrapper *aw) { _PyObject_GC_UNTRACK((PyObject *)aw); - Py_CLEAR(aw->aw_aiter); + Py_CLEAR(aw->ags_aiter); PyObject_GC_Del(aw); } @@ -1081,7 +1211,817 @@ _PyAIterWrapper_New(PyObject *aiter) return NULL; } Py_INCREF(aiter); - aw->aw_aiter = aiter; + aw->ags_aiter = aiter; _PyObject_GC_TRACK(aw); return (PyObject *)aw; } + + +/* ========= Asynchronous Generators ========= */ + + +typedef enum { + AWAITABLE_STATE_INIT, /* new awaitable, has not yet been iterated */ + AWAITABLE_STATE_ITER, /* being iterated */ + AWAITABLE_STATE_CLOSED, /* closed */ +} AwaitableState; + + +typedef struct { + PyObject_HEAD + PyAsyncGenObject *ags_gen; + + /* Can be NULL, when in the __anext__() mode + (equivalent of "asend(None)") */ + PyObject *ags_sendval; + + AwaitableState ags_state; +} PyAsyncGenASend; + + +typedef struct { + PyObject_HEAD + PyAsyncGenObject *agt_gen; + + /* Can be NULL, when in the "aclose()" mode + (equivalent of "athrow(GeneratorExit)") */ + PyObject *agt_args; + + AwaitableState agt_state; +} PyAsyncGenAThrow; + + +typedef struct { + PyObject_HEAD + PyObject *agw_val; +} _PyAsyncGenWrappedValue; + + +#ifndef _PyAsyncGen_MAXFREELIST +#define _PyAsyncGen_MAXFREELIST 80 +#endif + +/* Freelists boost performance 6-10%; they also reduce memory + fragmentation, as _PyAsyncGenWrappedValue and PyAsyncGenASend + are short-living objects that are instantiated for every + __anext__ call. +*/ + +static _PyAsyncGenWrappedValue *ag_value_freelist[_PyAsyncGen_MAXFREELIST]; +static int ag_value_freelist_free = 0; + +static PyAsyncGenASend *ag_asend_freelist[_PyAsyncGen_MAXFREELIST]; +static int ag_asend_freelist_free = 0; + +#define _PyAsyncGenWrappedValue_CheckExact(o) \ + (Py_TYPE(o) == &_PyAsyncGenWrappedValue_Type) + +#define PyAsyncGenASend_CheckExact(o) \ + (Py_TYPE(o) == &_PyAsyncGenASend_Type) + + +static int +async_gen_traverse(PyAsyncGenObject *gen, visitproc visit, void *arg) +{ + Py_VISIT(gen->ag_finalizer); + return gen_traverse((PyGenObject*)gen, visit, arg); +} + + +static PyObject * +async_gen_repr(PyAsyncGenObject *o) +{ + return PyUnicode_FromFormat("", + o->ag_qualname, o); +} + + +static int +async_gen_init_hooks(PyAsyncGenObject *o) +{ + PyThreadState *tstate; + PyObject *finalizer; + PyObject *firstiter; + + if (o->ag_hooks_inited) { + return 0; + } + + o->ag_hooks_inited = 1; + + tstate = PyThreadState_GET(); + + finalizer = tstate->async_gen_finalizer; + if (finalizer) { + Py_INCREF(finalizer); + o->ag_finalizer = finalizer; + } + + firstiter = tstate->async_gen_firstiter; + if (firstiter) { + PyObject *res; + + Py_INCREF(firstiter); + res = PyObject_CallFunction(firstiter, "O", o); + Py_DECREF(firstiter); + if (res == NULL) { + return 1; + } + Py_DECREF(res); + } + + return 0; +} + + +static PyObject * +async_gen_anext(PyAsyncGenObject *o) +{ + if (async_gen_init_hooks(o)) { + return NULL; + } + return async_gen_asend_new(o, NULL); +} + + +static PyObject * +async_gen_asend(PyAsyncGenObject *o, PyObject *arg) +{ + if (async_gen_init_hooks(o)) { + return NULL; + } + return async_gen_asend_new(o, arg); +} + + +static PyObject * +async_gen_aclose(PyAsyncGenObject *o, PyObject *arg) +{ + if (async_gen_init_hooks(o)) { + return NULL; + } + return async_gen_athrow_new(o, NULL); +} + +static PyObject * +async_gen_athrow(PyAsyncGenObject *o, PyObject *args) +{ + if (async_gen_init_hooks(o)) { + return NULL; + } + return async_gen_athrow_new(o, args); +} + + +static PyGetSetDef async_gen_getsetlist[] = { + {"__name__", (getter)gen_get_name, (setter)gen_set_name, + PyDoc_STR("name of the async generator")}, + {"__qualname__", (getter)gen_get_qualname, (setter)gen_set_qualname, + PyDoc_STR("qualified name of the async generator")}, + {"ag_await", (getter)coro_get_cr_await, NULL, + PyDoc_STR("object being awaited on, or None")}, + {NULL} /* Sentinel */ +}; + +static PyMemberDef async_gen_memberlist[] = { + {"ag_frame", T_OBJECT, offsetof(PyAsyncGenObject, ag_frame), READONLY}, + {"ag_running", T_BOOL, offsetof(PyAsyncGenObject, ag_running), READONLY}, + {"ag_code", T_OBJECT, offsetof(PyAsyncGenObject, ag_code), READONLY}, + {NULL} /* Sentinel */ +}; + +PyDoc_STRVAR(async_aclose_doc, +"aclose() -> raise GeneratorExit inside generator."); + +PyDoc_STRVAR(async_asend_doc, +"asend(v) -> send 'v' in generator."); + +PyDoc_STRVAR(async_athrow_doc, +"athrow(typ[,val[,tb]]) -> raise exception in generator."); + +static PyMethodDef async_gen_methods[] = { + {"asend", (PyCFunction)async_gen_asend, METH_O, async_asend_doc}, + {"athrow",(PyCFunction)async_gen_athrow, METH_VARARGS, async_athrow_doc}, + {"aclose", (PyCFunction)async_gen_aclose, METH_NOARGS, async_aclose_doc}, + {NULL, NULL} /* Sentinel */ +}; + + +static PyAsyncMethods async_gen_as_async = { + 0, /* am_await */ + PyObject_SelfIter, /* am_aiter */ + (unaryfunc)async_gen_anext /* am_anext */ +}; + + +PyTypeObject PyAsyncGen_Type = { + PyVarObject_HEAD_INIT(&PyType_Type, 0) + "async_generator", /* tp_name */ + sizeof(PyAsyncGenObject), /* tp_basicsize */ + 0, /* tp_itemsize */ + /* methods */ + (destructor)gen_dealloc, /* tp_dealloc */ + 0, /* tp_print */ + 0, /* tp_getattr */ + 0, /* tp_setattr */ + &async_gen_as_async, /* tp_as_async */ + (reprfunc)async_gen_repr, /* tp_repr */ + 0, /* tp_as_number */ + 0, /* tp_as_sequence */ + 0, /* tp_as_mapping */ + 0, /* tp_hash */ + 0, /* tp_call */ + 0, /* tp_str */ + PyObject_GenericGetAttr, /* tp_getattro */ + 0, /* tp_setattro */ + 0, /* tp_as_buffer */ + Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | + Py_TPFLAGS_HAVE_FINALIZE, /* tp_flags */ + 0, /* tp_doc */ + (traverseproc)async_gen_traverse, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + offsetof(PyAsyncGenObject, ag_weakreflist), /* tp_weaklistoffset */ + 0, /* tp_iter */ + 0, /* tp_iternext */ + async_gen_methods, /* tp_methods */ + async_gen_memberlist, /* tp_members */ + async_gen_getsetlist, /* tp_getset */ + 0, /* tp_base */ + 0, /* tp_dict */ + 0, /* tp_descr_get */ + 0, /* tp_descr_set */ + 0, /* tp_dictoffset */ + 0, /* tp_init */ + 0, /* tp_alloc */ + 0, /* tp_new */ + 0, /* tp_free */ + 0, /* tp_is_gc */ + 0, /* tp_bases */ + 0, /* tp_mro */ + 0, /* tp_cache */ + 0, /* tp_subclasses */ + 0, /* tp_weaklist */ + 0, /* tp_del */ + 0, /* tp_version_tag */ + _PyGen_Finalize, /* tp_finalize */ +}; + + +PyObject * +PyAsyncGen_New(PyFrameObject *f, PyObject *name, PyObject *qualname) +{ + PyAsyncGenObject *o; + o = (PyAsyncGenObject *)gen_new_with_qualname( + &PyAsyncGen_Type, f, name, qualname); + if (o == NULL) { + return NULL; + } + o->ag_finalizer = NULL; + o->ag_closed = 0; + o->ag_hooks_inited = 0; + return (PyObject*)o; +} + + +int +PyAsyncGen_ClearFreeLists(void) +{ + int ret = ag_value_freelist_free + ag_asend_freelist_free; + + while (ag_value_freelist_free) { + _PyAsyncGenWrappedValue *o; + o = ag_value_freelist[--ag_value_freelist_free]; + assert(_PyAsyncGenWrappedValue_CheckExact(o)); + PyObject_Del(o); + } + + while (ag_asend_freelist_free) { + PyAsyncGenASend *o; + o = ag_asend_freelist[--ag_asend_freelist_free]; + assert(Py_TYPE(o) == &_PyAsyncGenASend_Type); + PyObject_Del(o); + } + + return ret; +} + +void +PyAsyncGen_Fini(void) +{ + PyAsyncGen_ClearFreeLists(); +} + + +static PyObject * +async_gen_unwrap_value(PyAsyncGenObject *gen, PyObject *result) +{ + if (result == NULL) { + if (!PyErr_Occurred()) { + PyErr_SetNone(PyExc_StopAsyncIteration); + } + + if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration) + || PyErr_ExceptionMatches(PyExc_GeneratorExit) + ) { + gen->ag_closed = 1; + } + + return NULL; + } + + if (_PyAsyncGenWrappedValue_CheckExact(result)) { + /* async yield */ + PyObject *e = PyObject_CallFunctionObjArgs( + PyExc_StopIteration, + ((_PyAsyncGenWrappedValue*)result)->agw_val, + NULL); + Py_DECREF(result); + if (e == NULL) { + return NULL; + } + PyErr_SetObject(PyExc_StopIteration, e); + Py_DECREF(e); + return NULL; + } + + return result; +} + + +/* ---------- Async Generator ASend Awaitable ------------ */ + + +static void +async_gen_asend_dealloc(PyAsyncGenASend *o) +{ + Py_CLEAR(o->ags_gen); + Py_CLEAR(o->ags_sendval); + if (ag_asend_freelist_free < _PyAsyncGen_MAXFREELIST) { + assert(PyAsyncGenASend_CheckExact(o)); + ag_asend_freelist[ag_asend_freelist_free++] = o; + } else { + PyObject_Del(o); + } +} + + +static PyObject * +async_gen_asend_send(PyAsyncGenASend *o, PyObject *arg) +{ + PyObject *result; + + if (o->ags_state == AWAITABLE_STATE_CLOSED) { + PyErr_SetNone(PyExc_StopIteration); + return NULL; + } + + if (o->ags_state == AWAITABLE_STATE_INIT) { + if (arg == NULL || arg == Py_None) { + arg = o->ags_sendval; + } + o->ags_state = AWAITABLE_STATE_ITER; + } + + result = gen_send_ex((PyGenObject*)o->ags_gen, arg, 0, 0); + result = async_gen_unwrap_value(o->ags_gen, result); + + if (result == NULL) { + o->ags_state = AWAITABLE_STATE_CLOSED; + } + + return result; +} + + +static PyObject * +async_gen_asend_iternext(PyAsyncGenASend *o) +{ + return async_gen_asend_send(o, NULL); +} + + +static PyObject * +async_gen_asend_throw(PyAsyncGenASend *o, PyObject *args) +{ + PyObject *result; + + if (o->ags_state == AWAITABLE_STATE_CLOSED) { + PyErr_SetNone(PyExc_StopIteration); + return NULL; + } + + result = gen_throw((PyGenObject*)o->ags_gen, args); + result = async_gen_unwrap_value(o->ags_gen, result); + + if (result == NULL) { + o->ags_state = AWAITABLE_STATE_CLOSED; + } + + return result; +} + + +static PyObject * +async_gen_asend_close(PyAsyncGenASend *o, PyObject *args) +{ + o->ags_state = AWAITABLE_STATE_CLOSED; + Py_RETURN_NONE; +} + + +static PyMethodDef async_gen_asend_methods[] = { + {"send", (PyCFunction)async_gen_asend_send, METH_O, send_doc}, + {"throw", (PyCFunction)async_gen_asend_throw, METH_VARARGS, throw_doc}, + {"close", (PyCFunction)async_gen_asend_close, METH_NOARGS, close_doc}, + {NULL, NULL} /* Sentinel */ +}; + + +static PyAsyncMethods async_gen_asend_as_async = { + PyObject_SelfIter, /* am_await */ + 0, /* am_aiter */ + 0 /* am_anext */ +}; + + +PyTypeObject _PyAsyncGenASend_Type = { + PyVarObject_HEAD_INIT(&PyType_Type, 0) + "async_generator_asend", /* tp_name */ + sizeof(PyAsyncGenASend), /* tp_basicsize */ + 0, /* tp_itemsize */ + /* methods */ + (destructor)async_gen_asend_dealloc, /* tp_dealloc */ + 0, /* tp_print */ + 0, /* tp_getattr */ + 0, /* tp_setattr */ + &async_gen_asend_as_async, /* tp_as_async */ + 0, /* tp_repr */ + 0, /* tp_as_number */ + 0, /* tp_as_sequence */ + 0, /* tp_as_mapping */ + 0, /* tp_hash */ + 0, /* tp_call */ + 0, /* tp_str */ + PyObject_GenericGetAttr, /* tp_getattro */ + 0, /* tp_setattro */ + 0, /* tp_as_buffer */ + Py_TPFLAGS_DEFAULT, /* tp_flags */ + 0, /* tp_doc */ + 0, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + 0, /* tp_weaklistoffset */ + PyObject_SelfIter, /* tp_iter */ + (iternextfunc)async_gen_asend_iternext, /* tp_iternext */ + async_gen_asend_methods, /* tp_methods */ + 0, /* tp_members */ + 0, /* tp_getset */ + 0, /* tp_base */ + 0, /* tp_dict */ + 0, /* tp_descr_get */ + 0, /* tp_descr_set */ + 0, /* tp_dictoffset */ + 0, /* tp_init */ + 0, /* tp_alloc */ + 0, /* tp_new */ +}; + + +static PyObject * +async_gen_asend_new(PyAsyncGenObject *gen, PyObject *sendval) +{ + PyAsyncGenASend *o; + if (ag_asend_freelist_free) { + ag_asend_freelist_free--; + o = ag_asend_freelist[ag_asend_freelist_free]; + _Py_NewReference((PyObject *)o); + } else { + o = PyObject_New(PyAsyncGenASend, &_PyAsyncGenASend_Type); + if (o == NULL) { + return NULL; + } + } + + Py_INCREF(gen); + o->ags_gen = gen; + + Py_XINCREF(sendval); + o->ags_sendval = sendval; + + o->ags_state = AWAITABLE_STATE_INIT; + return (PyObject*)o; +} + + +/* ---------- Async Generator Value Wrapper ------------ */ + + +static void +async_gen_wrapped_val_dealloc(_PyAsyncGenWrappedValue *o) +{ + Py_CLEAR(o->agw_val); + if (ag_value_freelist_free < _PyAsyncGen_MAXFREELIST) { + assert(_PyAsyncGenWrappedValue_CheckExact(o)); + ag_value_freelist[ag_value_freelist_free++] = o; + } else { + PyObject_Del(o); + } +} + + +PyTypeObject _PyAsyncGenWrappedValue_Type = { + PyVarObject_HEAD_INIT(&PyType_Type, 0) + "async_generator_wrapped_value", /* tp_name */ + sizeof(_PyAsyncGenWrappedValue), /* tp_basicsize */ + 0, /* tp_itemsize */ + /* methods */ + (destructor)async_gen_wrapped_val_dealloc, /* tp_dealloc */ + 0, /* tp_print */ + 0, /* tp_getattr */ + 0, /* tp_setattr */ + 0, /* tp_as_async */ + 0, /* tp_repr */ + 0, /* tp_as_number */ + 0, /* tp_as_sequence */ + 0, /* tp_as_mapping */ + 0, /* tp_hash */ + 0, /* tp_call */ + 0, /* tp_str */ + PyObject_GenericGetAttr, /* tp_getattro */ + 0, /* tp_setattro */ + 0, /* tp_as_buffer */ + Py_TPFLAGS_DEFAULT, /* tp_flags */ + 0, /* tp_doc */ + 0, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + 0, /* tp_weaklistoffset */ + 0, /* tp_iter */ + 0, /* tp_iternext */ + 0, /* tp_methods */ + 0, /* tp_members */ + 0, /* tp_getset */ + 0, /* tp_base */ + 0, /* tp_dict */ + 0, /* tp_descr_get */ + 0, /* tp_descr_set */ + 0, /* tp_dictoffset */ + 0, /* tp_init */ + 0, /* tp_alloc */ + 0, /* tp_new */ +}; + + +PyObject * +_PyAsyncGenValueWrapperNew(PyObject *val) +{ + _PyAsyncGenWrappedValue *o; + assert(val); + + if (ag_value_freelist_free) { + ag_value_freelist_free--; + o = ag_value_freelist[ag_value_freelist_free]; + assert(_PyAsyncGenWrappedValue_CheckExact(o)); + _Py_NewReference((PyObject*)o); + } else { + o = PyObject_New(_PyAsyncGenWrappedValue, &_PyAsyncGenWrappedValue_Type); + if (o == NULL) { + return NULL; + } + } + o->agw_val = val; + Py_INCREF(val); + return (PyObject*)o; +} + + +/* ---------- Async Generator AThrow awaitable ------------ */ + + +static void +async_gen_athrow_dealloc(PyAsyncGenAThrow *o) +{ + Py_CLEAR(o->agt_gen); + Py_CLEAR(o->agt_args); + PyObject_Del(o); +} + + +static PyObject * +async_gen_athrow_send(PyAsyncGenAThrow *o, PyObject *arg) +{ + PyGenObject *gen = (PyGenObject*)o->agt_gen; + PyFrameObject *f = gen->gi_frame; + PyObject *retval; + + if (f == NULL || f->f_stacktop == NULL || + o->agt_state == AWAITABLE_STATE_CLOSED) { + PyErr_SetNone(PyExc_StopIteration); + return NULL; + } + + if (o->agt_state == AWAITABLE_STATE_INIT) { + if (o->agt_gen->ag_closed) { + PyErr_SetNone(PyExc_StopIteration); + return NULL; + } + + if (arg != Py_None) { + PyErr_SetString(PyExc_RuntimeError, NON_INIT_CORO_MSG); + return NULL; + } + + o->agt_state = AWAITABLE_STATE_ITER; + + if (o->agt_args == NULL) { + /* aclose() mode */ + o->agt_gen->ag_closed = 1; + + retval = _gen_throw((PyGenObject *)gen, + 0, /* Do not close generator when + PyExc_GeneratorExit is passed */ + PyExc_GeneratorExit, NULL, NULL); + + if (retval && _PyAsyncGenWrappedValue_CheckExact(retval)) { + Py_DECREF(retval); + goto yield_close; + } + } else { + PyObject *typ; + PyObject *tb = NULL; + PyObject *val = NULL; + + if (!PyArg_UnpackTuple(o->agt_args, "athrow", 1, 3, + &typ, &val, &tb)) { + return NULL; + } + + retval = _gen_throw((PyGenObject *)gen, + 0, /* Do not close generator when + PyExc_GeneratorExit is passed */ + typ, val, tb); + retval = async_gen_unwrap_value(o->agt_gen, retval); + } + if (retval == NULL) { + goto check_error; + } + return retval; + } + + assert(o->agt_state == AWAITABLE_STATE_ITER); + + retval = gen_send_ex((PyGenObject *)gen, arg, 0, 0); + if (o->agt_args) { + return async_gen_unwrap_value(o->agt_gen, retval); + } else { + /* aclose() mode */ + if (retval) { + if (_PyAsyncGenWrappedValue_CheckExact(retval)) { + Py_DECREF(retval); + goto yield_close; + } + else { + return retval; + } + } + else { + goto check_error; + } + } + +yield_close: + PyErr_SetString( + PyExc_RuntimeError, ASYNC_GEN_IGNORED_EXIT_MSG); + return NULL; + +check_error: + if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration) + || PyErr_ExceptionMatches(PyExc_GeneratorExit) + ) { + o->agt_state = AWAITABLE_STATE_CLOSED; + PyErr_Clear(); /* ignore these errors */ + PyErr_SetNone(PyExc_StopIteration); + } + return NULL; +} + + +static PyObject * +async_gen_athrow_throw(PyAsyncGenAThrow *o, PyObject *args) +{ + PyObject *retval; + + if (o->agt_state == AWAITABLE_STATE_INIT) { + PyErr_SetString(PyExc_RuntimeError, NON_INIT_CORO_MSG); + return NULL; + } + + if (o->agt_state == AWAITABLE_STATE_CLOSED) { + PyErr_SetNone(PyExc_StopIteration); + return NULL; + } + + retval = gen_throw((PyGenObject*)o->agt_gen, args); + if (o->agt_args) { + return async_gen_unwrap_value(o->agt_gen, retval); + } else { + /* aclose() mode */ + if (retval && _PyAsyncGenWrappedValue_CheckExact(retval)) { + Py_DECREF(retval); + PyErr_SetString(PyExc_RuntimeError, ASYNC_GEN_IGNORED_EXIT_MSG); + return NULL; + } + return retval; + } +} + + +static PyObject * +async_gen_athrow_iternext(PyAsyncGenAThrow *o) +{ + return async_gen_athrow_send(o, Py_None); +} + + +static PyObject * +async_gen_athrow_close(PyAsyncGenAThrow *o, PyObject *args) +{ + o->agt_state = AWAITABLE_STATE_CLOSED; + Py_RETURN_NONE; +} + + +static PyMethodDef async_gen_athrow_methods[] = { + {"send", (PyCFunction)async_gen_athrow_send, METH_O, send_doc}, + {"throw", (PyCFunction)async_gen_athrow_throw, METH_VARARGS, throw_doc}, + {"close", (PyCFunction)async_gen_athrow_close, METH_NOARGS, close_doc}, + {NULL, NULL} /* Sentinel */ +}; + + +static PyAsyncMethods async_gen_athrow_as_async = { + PyObject_SelfIter, /* am_await */ + 0, /* am_aiter */ + 0 /* am_anext */ +}; + + +PyTypeObject _PyAsyncGenAThrow_Type = { + PyVarObject_HEAD_INIT(&PyType_Type, 0) + "async_generator_athrow", /* tp_name */ + sizeof(PyAsyncGenAThrow), /* tp_basicsize */ + 0, /* tp_itemsize */ + /* methods */ + (destructor)async_gen_athrow_dealloc, /* tp_dealloc */ + 0, /* tp_print */ + 0, /* tp_getattr */ + 0, /* tp_setattr */ + &async_gen_athrow_as_async, /* tp_as_async */ + 0, /* tp_repr */ + 0, /* tp_as_number */ + 0, /* tp_as_sequence */ + 0, /* tp_as_mapping */ + 0, /* tp_hash */ + 0, /* tp_call */ + 0, /* tp_str */ + PyObject_GenericGetAttr, /* tp_getattro */ + 0, /* tp_setattro */ + 0, /* tp_as_buffer */ + Py_TPFLAGS_DEFAULT, /* tp_flags */ + 0, /* tp_doc */ + 0, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + 0, /* tp_weaklistoffset */ + PyObject_SelfIter, /* tp_iter */ + (iternextfunc)async_gen_athrow_iternext, /* tp_iternext */ + async_gen_athrow_methods, /* tp_methods */ + 0, /* tp_members */ + 0, /* tp_getset */ + 0, /* tp_base */ + 0, /* tp_dict */ + 0, /* tp_descr_get */ + 0, /* tp_descr_set */ + 0, /* tp_dictoffset */ + 0, /* tp_init */ + 0, /* tp_alloc */ + 0, /* tp_new */ +}; + + +static PyObject * +async_gen_athrow_new(PyAsyncGenObject *gen, PyObject *args) +{ + PyAsyncGenAThrow *o; + o = PyObject_New(PyAsyncGenAThrow, &_PyAsyncGenAThrow_Type); + if (o == NULL) { + return NULL; + } + o->agt_gen = gen; + o->agt_args = args; + o->agt_state = AWAITABLE_STATE_INIT; + Py_INCREF(gen); + Py_XINCREF(args); + return (PyObject*)o; +} diff --git a/Python/ceval.c b/Python/ceval.c index a52ee8a..f737a2f 100644 --- a/Python/ceval.c +++ b/Python/ceval.c @@ -1204,7 +1204,7 @@ _PyEval_EvalFrameDefault(PyFrameObject *f, int throwflag) f->f_stacktop = NULL; /* remains NULL unless yield suspends frame */ f->f_executing = 1; - if (co->co_flags & (CO_GENERATOR | CO_COROUTINE)) { + if (co->co_flags & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR)) { if (!throwflag && f->f_exc_type != NULL && f->f_exc_type != Py_None) { /* We were in an except handler when we left, restore the exception state which was put aside @@ -2083,36 +2083,45 @@ _PyEval_EvalFrameDefault(PyFrameObject *f, int throwflag) PyObject *aiter = TOP(); PyTypeObject *type = Py_TYPE(aiter); - if (type->tp_as_async != NULL) - getter = type->tp_as_async->am_anext; + if (PyAsyncGen_CheckExact(aiter)) { + awaitable = type->tp_as_async->am_anext(aiter); + if (awaitable == NULL) { + goto error; + } + } else { + if (type->tp_as_async != NULL){ + getter = type->tp_as_async->am_anext; + } - if (getter != NULL) { - next_iter = (*getter)(aiter); - if (next_iter == NULL) { + if (getter != NULL) { + next_iter = (*getter)(aiter); + if (next_iter == NULL) { + goto error; + } + } + else { + PyErr_Format( + PyExc_TypeError, + "'async for' requires an iterator with " + "__anext__ method, got %.100s", + type->tp_name); goto error; } - } - else { - PyErr_Format( - PyExc_TypeError, - "'async for' requires an iterator with " - "__anext__ method, got %.100s", - type->tp_name); - goto error; - } - awaitable = _PyCoro_GetAwaitableIter(next_iter); - if (awaitable == NULL) { - PyErr_Format( - PyExc_TypeError, - "'async for' received an invalid object " - "from __anext__: %.100s", - Py_TYPE(next_iter)->tp_name); + awaitable = _PyCoro_GetAwaitableIter(next_iter); + if (awaitable == NULL) { + PyErr_Format( + PyExc_TypeError, + "'async for' received an invalid object " + "from __anext__: %.100s", + Py_TYPE(next_iter)->tp_name); - Py_DECREF(next_iter); - goto error; - } else - Py_DECREF(next_iter); + Py_DECREF(next_iter); + goto error; + } else { + Py_DECREF(next_iter); + } + } PUSH(awaitable); PREDICT(LOAD_CONST); @@ -2187,6 +2196,17 @@ _PyEval_EvalFrameDefault(PyFrameObject *f, int throwflag) TARGET(YIELD_VALUE) { retval = POP(); + + if (co->co_flags & CO_ASYNC_GENERATOR) { + PyObject *w = _PyAsyncGenValueWrapperNew(retval); + Py_DECREF(retval); + if (w == NULL) { + retval = NULL; + goto error; + } + retval = w; + } + f->f_stacktop = stack_pointer; why = WHY_YIELD; goto fast_yield; @@ -3712,7 +3732,7 @@ fast_block_end: assert((retval != NULL) ^ (PyErr_Occurred() != NULL)); fast_yield: - if (co->co_flags & (CO_GENERATOR | CO_COROUTINE)) { + if (co->co_flags & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR)) { /* The purpose of this block is to put aside the generator's exception state and restore that of the calling frame. If the current @@ -4156,8 +4176,8 @@ _PyEval_EvalCodeWithName(PyObject *_co, PyObject *globals, PyObject *locals, freevars[PyTuple_GET_SIZE(co->co_cellvars) + i] = o; } - /* Handle generator/coroutine */ - if (co->co_flags & (CO_GENERATOR | CO_COROUTINE)) { + /* Handle generator/coroutine/asynchronous generator */ + if (co->co_flags & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR)) { PyObject *gen; PyObject *coro_wrapper = tstate->coroutine_wrapper; int is_coro = co->co_flags & CO_COROUTINE; @@ -4182,6 +4202,8 @@ _PyEval_EvalCodeWithName(PyObject *_co, PyObject *globals, PyObject *locals, * and return that as the value. */ if (is_coro) { gen = PyCoro_New(f, name, qualname); + } else if (co->co_flags & CO_ASYNC_GENERATOR) { + gen = PyAsyncGen_New(f, name, qualname); } else { gen = PyGen_NewWithQualName(f, name, qualname); } @@ -4660,6 +4682,38 @@ _PyEval_GetCoroutineWrapper(void) return tstate->coroutine_wrapper; } +void +_PyEval_SetAsyncGenFirstiter(PyObject *firstiter) +{ + PyThreadState *tstate = PyThreadState_GET(); + + Py_XINCREF(firstiter); + Py_XSETREF(tstate->async_gen_firstiter, firstiter); +} + +PyObject * +_PyEval_GetAsyncGenFirstiter(void) +{ + PyThreadState *tstate = PyThreadState_GET(); + return tstate->async_gen_firstiter; +} + +void +_PyEval_SetAsyncGenFinalizer(PyObject *finalizer) +{ + PyThreadState *tstate = PyThreadState_GET(); + + Py_XINCREF(finalizer); + Py_XSETREF(tstate->async_gen_finalizer, finalizer); +} + +PyObject * +_PyEval_GetAsyncGenFinalizer(void) +{ + PyThreadState *tstate = PyThreadState_GET(); + return tstate->async_gen_finalizer; +} + PyObject * PyEval_GetBuiltins(void) { diff --git a/Python/compile.c b/Python/compile.c index b46edd4..faae4f5 100644 --- a/Python/compile.c +++ b/Python/compile.c @@ -1886,8 +1886,6 @@ compiler_function(struct compiler *c, stmt_ty s, int is_async) return 0; } - if (is_async) - co->co_flags |= CO_COROUTINE; compiler_make_closure(c, co, funcflags, qualname); Py_DECREF(qualname); Py_DECREF(co); @@ -2801,6 +2799,9 @@ compiler_visit_stmt(struct compiler *c, stmt_ty s) if (c->u->u_ste->ste_type != FunctionBlock) return compiler_error(c, "'return' outside function"); if (s->v.Return.value) { + if (c->u->u_ste->ste_coroutine && c->u->u_ste->ste_generator) + return compiler_error( + c, "'return' with value in async generator"); VISIT(c, expr, s->v.Return.value); } else @@ -4115,8 +4116,6 @@ compiler_visit_expr(struct compiler *c, expr_ty e) case Yield_kind: if (c->u->u_ste->ste_type != FunctionBlock) return compiler_error(c, "'yield' outside function"); - if (c->u->u_scope_type == COMPILER_SCOPE_ASYNC_FUNCTION) - return compiler_error(c, "'yield' inside async function"); if (e->v.Yield.value) { VISIT(c, expr, e->v.Yield.value); } @@ -4992,8 +4991,12 @@ compute_code_flags(struct compiler *c) flags |= CO_NEWLOCALS | CO_OPTIMIZED; if (ste->ste_nested) flags |= CO_NESTED; - if (ste->ste_generator) + if (ste->ste_generator && !ste->ste_coroutine) flags |= CO_GENERATOR; + if (!ste->ste_generator && ste->ste_coroutine) + flags |= CO_COROUTINE; + if (ste->ste_generator && ste->ste_coroutine) + flags |= CO_ASYNC_GENERATOR; if (ste->ste_varargs) flags |= CO_VARARGS; if (ste->ste_varkeywords) diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c index a2399ed..f93afd2 100644 --- a/Python/pylifecycle.c +++ b/Python/pylifecycle.c @@ -694,6 +694,7 @@ Py_FinalizeEx(void) _PyGC_Fini(); _PyRandom_Fini(); _PyArg_Fini(); + PyAsyncGen_Fini(); /* Cleanup Unicode implementation */ _PyUnicode_Fini(); diff --git a/Python/pystate.c b/Python/pystate.c index 959354d..a0a8c97 100644 --- a/Python/pystate.c +++ b/Python/pystate.c @@ -229,6 +229,9 @@ new_threadstate(PyInterpreterState *interp, int init) tstate->in_coroutine_wrapper = 0; tstate->co_extra_user_count = 0; + tstate->async_gen_firstiter = NULL; + tstate->async_gen_finalizer = NULL; + if (init) _PyThreadState_Init(tstate); @@ -408,6 +411,8 @@ PyThreadState_Clear(PyThreadState *tstate) Py_CLEAR(tstate->c_traceobj); Py_CLEAR(tstate->coroutine_wrapper); + Py_CLEAR(tstate->async_gen_firstiter); + Py_CLEAR(tstate->async_gen_finalizer); } diff --git a/Python/symtable.c b/Python/symtable.c index b8d9398..e55813f 100644 --- a/Python/symtable.c +++ b/Python/symtable.c @@ -63,6 +63,7 @@ ste_new(struct symtable *st, identifier name, _Py_block_ty block, ste->ste_nested = 1; ste->ste_child_free = 0; ste->ste_generator = 0; + ste->ste_coroutine = 0; ste->ste_returns_value = 0; ste->ste_needs_class_closure = 0; @@ -378,7 +379,7 @@ error_at_directive(PySTEntryObject *ste, PyObject *name) PyLong_AsLong(PyTuple_GET_ITEM(data, 2))); return 0; - } + } } PyErr_SetString(PyExc_RuntimeError, "BUG: internal directive bookkeeping broken"); @@ -1397,6 +1398,7 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s) FunctionBlock, (void *)s, s->lineno, s->col_offset)) VISIT_QUIT(st, 0); + st->st_cur->ste_coroutine = 1; VISIT(st, arguments, s->v.AsyncFunctionDef.args); VISIT_SEQ(st, stmt, s->v.AsyncFunctionDef.body); if (!symtable_exit_block(st, s)) @@ -1492,7 +1494,7 @@ symtable_visit_expr(struct symtable *st, expr_ty e) break; case Await_kind: VISIT(st, expr, e->v.Await.value); - st->st_cur->ste_generator = 1; + st->st_cur->ste_coroutine = 1; break; case Compare_kind: VISIT(st, expr, e->v.Compare.left); diff --git a/Python/sysmodule.c b/Python/sysmodule.c index 0fe76b7..3a02ae9 100644 --- a/Python/sysmodule.c +++ b/Python/sysmodule.c @@ -717,6 +717,113 @@ Return the wrapper for coroutine objects set by sys.set_coroutine_wrapper." ); +static PyTypeObject AsyncGenHooksType; + +PyDoc_STRVAR(asyncgen_hooks_doc, +"asyncgen_hooks\n\ +\n\ +A struct sequence providing information about asynhronous\n\ +generators hooks. The attributes are read only."); + +static PyStructSequence_Field asyncgen_hooks_fields[] = { + {"firstiter", "Hook to intercept first iteration"}, + {"finalizer", "Hook to intercept finalization"}, + {0} +}; + +static PyStructSequence_Desc asyncgen_hooks_desc = { + "asyncgen_hooks", /* name */ + asyncgen_hooks_doc, /* doc */ + asyncgen_hooks_fields , /* fields */ + 2 +}; + + +static PyObject * +sys_set_asyncgen_hooks(PyObject *self, PyObject *args, PyObject *kw) +{ + static char *keywords[] = {"firstiter", "finalizer", NULL}; + PyObject *firstiter = NULL; + PyObject *finalizer = NULL; + + if (!PyArg_ParseTupleAndKeywords( + args, kw, "|OO", keywords, + &firstiter, &finalizer)) { + return NULL; + } + + if (finalizer && finalizer != Py_None) { + if (!PyCallable_Check(finalizer)) { + PyErr_Format(PyExc_TypeError, + "callable finalizer expected, got %.50s", + Py_TYPE(finalizer)->tp_name); + return NULL; + } + _PyEval_SetAsyncGenFinalizer(finalizer); + } + else if (finalizer == Py_None) { + _PyEval_SetAsyncGenFinalizer(NULL); + } + + if (firstiter && firstiter != Py_None) { + if (!PyCallable_Check(firstiter)) { + PyErr_Format(PyExc_TypeError, + "callable firstiter expected, got %.50s", + Py_TYPE(firstiter)->tp_name); + return NULL; + } + _PyEval_SetAsyncGenFirstiter(firstiter); + } + else if (firstiter == Py_None) { + _PyEval_SetAsyncGenFirstiter(NULL); + } + + Py_RETURN_NONE; +} + +PyDoc_STRVAR(set_asyncgen_hooks_doc, +"set_asyncgen_hooks(*, firstiter=None, finalizer=None)\n\ +\n\ +Set a finalizer for async generators objects." +); + +static PyObject * +sys_get_asyncgen_hooks(PyObject *self, PyObject *args) +{ + PyObject *res; + PyObject *firstiter = _PyEval_GetAsyncGenFirstiter(); + PyObject *finalizer = _PyEval_GetAsyncGenFinalizer(); + + res = PyStructSequence_New(&AsyncGenHooksType); + if (res == NULL) { + return NULL; + } + + if (firstiter == NULL) { + firstiter = Py_None; + } + + if (finalizer == NULL) { + finalizer = Py_None; + } + + Py_INCREF(firstiter); + PyStructSequence_SET_ITEM(res, 0, firstiter); + + Py_INCREF(finalizer); + PyStructSequence_SET_ITEM(res, 1, finalizer); + + return res; +} + +PyDoc_STRVAR(get_asyncgen_hooks_doc, +"get_asyncgen_hooks()\n\ +\n\ +Return a namedtuple of installed asynchronous generators hooks \ +(firstiter, finalizer)." +); + + static PyTypeObject Hash_InfoType; PyDoc_STRVAR(hash_info_doc, @@ -1315,6 +1422,10 @@ static PyMethodDef sys_methods[] = { set_coroutine_wrapper_doc}, {"get_coroutine_wrapper", sys_get_coroutine_wrapper, METH_NOARGS, get_coroutine_wrapper_doc}, + {"set_asyncgen_hooks", sys_set_asyncgen_hooks, + METH_VARARGS | METH_KEYWORDS, set_asyncgen_hooks_doc}, + {"get_asyncgen_hooks", sys_get_asyncgen_hooks, METH_NOARGS, + get_asyncgen_hooks_doc}, {NULL, NULL} /* sentinel */ }; @@ -1950,6 +2061,14 @@ _PySys_Init(void) SET_SYS_FROM_STRING("thread_info", PyThread_GetInfo()); #endif + /* initialize asyncgen_hooks */ + if (AsyncGenHooksType.tp_name == NULL) { + if (PyStructSequence_InitType2( + &AsyncGenHooksType, &asyncgen_hooks_desc) < 0) { + return NULL; + } + } + #undef SET_SYS_FROM_STRING #undef SET_SYS_FROM_STRING_BORROW if (PyErr_Occurred())