diff --git a/Lib/asyncio/events.py b/Lib/asyncio/events.py index 28a45fc3c..466db6d9a 100644 --- a/Lib/asyncio/events.py +++ b/Lib/asyncio/events.py @@ -6,7 +6,8 @@ __all__ = ['AbstractEventLoopPolicy', 'get_event_loop_policy', 'set_event_loop_policy', 'get_event_loop', 'set_event_loop', 'new_event_loop', 'get_child_watcher', 'set_child_watcher', - '_set_running_loop', '_get_running_loop', + '_set_running_loop', 'get_running_loop', + '_get_running_loop', ] import functools @@ -614,6 +615,18 @@ class _RunningLoop(threading.local): _running_loop = _RunningLoop() +def get_running_loop(): + """Return the running event loop. Raise a RuntimeError if there is none. + + This function is thread-specific. + """ + # NOTE: this function is implemented in C (see _asynciomodule.c) + loop = _get_running_loop() + if loop is None: + raise RuntimeError('no running event loop') + return loop + + def _get_running_loop(): """Return the running event loop or None. diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index d31c0109f..8c4f3cca7 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -1,6 +1,6 @@ """Support for tasks, coroutines and the scheduler.""" -__all__ = ['Task', +__all__ = ['Task', 'create_task', 'FIRST_COMPLETED', 'FIRST_EXCEPTION', 'ALL_COMPLETED', 'wait', 'wait_for', 'as_completed', 'sleep', 'async', 'gather', 'shield', 'ensure_future', 'run_coroutine_threadsafe', @@ -44,6 +44,16 @@ def _all_tasks_compat(loop=None): return {t for t in _all_tasks if futures._get_loop(t) is loop} +def _set_task_name(task, name): + if name is not None: + try: + set_name = task.set_name + except AttributeError: + pass + else: + set_name(name) + + class Task(futures.Future): """A coroutine wrapped in a Future.""" @@ -292,6 +302,17 @@ else: Task = _CTask = _asyncio.Task +def create_task(coro, *, name=None): + """Schedule the execution of a coroutine object in a spawn task. + + Return a Task object. + """ + loop = events.get_running_loop() + task = loop.create_task(coro) + _set_task_name(task, name) + return task + + # wait() and as_completed() similar to those in PEP 3148. FIRST_COMPLETED = concurrent.futures.FIRST_COMPLETED diff --git a/Lib/test/test_asyncgen.py b/Lib/test/test_asyncgen.py new file mode 100644 index 000000000..ab3743d08 --- /dev/null +++ b/Lib/test/test_asyncgen.py @@ -0,0 +1,1225 @@ +import inspect +import types +import unittest + +from test.support import import_module +asyncio = import_module("asyncio") + + +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: + an = g.__anext__() + try: + while True: + try: + an.__next__() + except StopIteration as ex: + if ex.args: + res.append(ex.args[0]) + break + else: + res.append('EMPTY StopIteration') + break + except StopAsyncIteration: + raise + except Exception as ex: + res.append(str(type(ex))) + break + except StopAsyncIteration: + res.append('STOP') + break + 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__() + + an = ai.__anext__() + self.assertEqual(an.__next__(), ('result',)) + + try: + an.__next__() + except StopIteration as ex: + self.assertEqual(ex.args[0], 123) + else: + self.fail('StopIteration was not raised') + + an = ai.__anext__() + self.assertEqual(an.__next__(), ('result',)) + + try: + an.__next__() + except StopAsyncIteration as ex: + self.assertFalse(ex.args) + else: + self.fail('StopAsyncIteration was not raised') + + @unittest.skip("TODO: RUSTPYTHON") + 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__() + an = ai.__anext__() + self.assertEqual(an.__next__(), ('result',)) + + try: + an.__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()) + + @unittest.skip("TODO: RUSTPYTHON") + 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()) + + @unittest.skip("TODO: RUSTPYTHON") + 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()) + + @unittest.skip("TODO: RUSTPYTHON") + 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_exception_11(self): + def sync_gen(): + yield 10 + yield 20 + + def sync_gen_wrapper(): + yield 1 + sg = sync_gen() + sg.send(None) + try: + sg.throw(GeneratorExit()) + except GeneratorExit: + yield 2 + yield 3 + + async def async_gen(): + yield 10 + yield 20 + + async def async_gen_wrapper(): + yield 1 + asg = async_gen() + await asg.asend(None) + try: + await asg.athrow(GeneratorExit()) + except GeneratorExit: + yield 2 + yield 3 + + self.compare_generators(sync_gen_wrapper(), async_gen_wrapper()) + + # TODO: RUSTPYTHON + @unittest.expectedFailure + 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 + asyncio.set_event_loop_policy(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) + yield 2 + await asyncio.sleep(0.01) + 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) + 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) + 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) + try: + yield 2 + yield 3 + except ZeroDivisionError: + yield 1000 + await asyncio.sleep(0.01) + 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_anext_06(self): + DONE = 0 + + # test synchronous generators + def foo(): + try: + yield + except: + pass + g = foo() + g.send(None) + with self.assertRaises(StopIteration): + g.send(None) + + # now with asynchronous generators + + async def gen(): + nonlocal DONE + try: + yield + except: + pass + DONE = 1 + + async def run(): + nonlocal DONE + g = gen() + await g.asend(None) + with self.assertRaises(StopAsyncIteration): + await g.asend(None) + DONE += 10 + + self.loop.run_until_complete(run()) + self.assertEqual(DONE, 11) + + def test_async_gen_asyncio_anext_tuple(self): + async def foo(): + try: + yield (1,) + except ZeroDivisionError: + yield (2,) + + async def run(): + it = foo().__aiter__() + + self.assertEqual(await it.__anext__(), (1,)) + with self.assertRaises(StopIteration) as cm: + it.__anext__().throw(ZeroDivisionError) + self.assertEqual(cm.exception.args[0], (2,)) + with self.assertRaises(StopAsyncIteration): + await it.__anext__() + + self.loop.run_until_complete(run()) + + def test_async_gen_asyncio_anext_stopiteration(self): + async def foo(): + try: + yield StopIteration(1) + except ZeroDivisionError: + yield StopIteration(3) + + async def run(): + it = foo().__aiter__() + + v = await it.__anext__() + self.assertIsInstance(v, StopIteration) + self.assertEqual(v.value, 1) + with self.assertRaises(StopIteration) as cm: + it.__anext__().throw(ZeroDivisionError) + v = cm.exception.args[0] + self.assertIsInstance(v, StopIteration) + self.assertEqual(v.value, 3) + 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) + 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) + await asyncio.sleep(0.01) + 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) + await asyncio.sleep(0.01) + DONE += 1 + DONE += 1000 + + async def run(): + gen = foo() + it = gen.__aiter__() + self.assertEqual(await it.__anext__(), 1) + await gen.aclose() + + self.loop.run_until_complete(run()) + self.assertEqual(DONE, 1) + + # Silence ResourceWarnings + fut.cancel() + self.loop.run_until_complete(asyncio.sleep(0.01)) + + # TODO: RUSTPYTHON + @unittest.expectedFailure + 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) + await asyncio.sleep(0.01) + DONE = 1 + + async def run(): + g = gen() + await g.__anext__() + await g.__anext__() + del g + + await asyncio.sleep(0.1) + + self.loop.run_until_complete(run()) + self.assertEqual(DONE, 1) + + def test_async_gen_asyncio_aclose_10(self): + DONE = 0 + + # test synchronous generators + def foo(): + try: + yield + except: + pass + g = foo() + g.send(None) + g.close() + + # now with asynchronous generators + + async def gen(): + nonlocal DONE + try: + yield + except: + pass + DONE = 1 + + async def run(): + nonlocal DONE + g = gen() + await g.asend(None) + await g.aclose() + DONE += 10 + + self.loop.run_until_complete(run()) + self.assertEqual(DONE, 11) + + def test_async_gen_asyncio_aclose_11(self): + DONE = 0 + + # test synchronous generators + def foo(): + try: + yield + except: + pass + yield + g = foo() + g.send(None) + with self.assertRaisesRegex(RuntimeError, 'ignored GeneratorExit'): + g.close() + + # now with asynchronous generators + + async def gen(): + nonlocal DONE + try: + yield + except: + pass + yield + DONE += 1 + + async def run(): + nonlocal DONE + g = gen() + await g.asend(None) + with self.assertRaisesRegex(RuntimeError, 'ignored GeneratorExit'): + await g.aclose() + DONE += 10 + + self.loop.run_until_complete(run()) + self.assertEqual(DONE, 10) + + # TODO: RUSTPYTHON + @unittest.expectedFailure + def test_async_gen_asyncio_aclose_12(self): + DONE = 0 + + async def target(): + await asyncio.sleep(0.01) + 1 / 0 + + async def foo(): + nonlocal DONE + task = asyncio.create_task(target()) + try: + yield 1 + finally: + try: + await task + except ZeroDivisionError: + DONE = 1 + + 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_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) + v = yield 1 + await asyncio.sleep(0.01) + yield v * 2 + await asyncio.sleep(0.01) + return + finally: + await asyncio.sleep(0.01) + await asyncio.sleep(0.01) + 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) + 1 / 0 + + async def gen(): + nonlocal DONE + try: + await asyncio.sleep(0.01) + v = yield 1 + await sleep_n_crash(0.01) + DONE += 1000 + yield v * 2 + finally: + await asyncio.sleep(0.01) + await asyncio.sleep(0.01) + 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) + + # TODO: RUSTPYTHON + @unittest.expectedFailure + 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) + self.loop.call_later(delay / 2, lambda: fut.cancel()) + return await fut + + async def gen(): + nonlocal DONE + try: + await asyncio.sleep(0.01) + v = yield 1 + await sleep_n_crash(0.01) + DONE += 1000 + yield v * 2 + finally: + await asyncio.sleep(0.01) + await asyncio.sleep(0.01) + 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) + try: + v = yield 1 + except FooEr: + v = 1000 + await asyncio.sleep(0.01) + yield v * 2 + await asyncio.sleep(0.01) + # return + finally: + await asyncio.sleep(0.01) + await asyncio.sleep(0.01) + 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) + + @unittest.skip("TODO: RUSTPYTHON") + 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) + self.loop.call_later(delay / 2, lambda: fut.cancel()) + return await fut + + async def gen(): + nonlocal DONE + try: + await asyncio.sleep(0.01) + try: + v = yield 1 + except FooEr: + await sleep_n_crash(0.01) + yield v * 2 + await asyncio.sleep(0.01) + # return + finally: + await asyncio.sleep(0.01) + await asyncio.sleep(0.01) + 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_athrow_03(self): + DONE = 0 + + # test synchronous generators + def foo(): + try: + yield + except: + pass + g = foo() + g.send(None) + with self.assertRaises(StopIteration): + g.throw(ValueError) + + # now with asynchronous generators + + async def gen(): + nonlocal DONE + try: + yield + except: + pass + DONE = 1 + + async def run(): + nonlocal DONE + g = gen() + await g.asend(None) + with self.assertRaises(StopAsyncIteration): + await g.athrow(ValueError) + DONE += 10 + + self.loop.run_until_complete(run()) + self.assertEqual(DONE, 11) + + def test_async_gen_asyncio_athrow_tuple(self): + async def gen(): + try: + yield 1 + except ZeroDivisionError: + yield (2,) + + async def run(): + g = gen() + v = await g.asend(None) + self.assertEqual(v, 1) + v = await g.athrow(ZeroDivisionError) + self.assertEqual(v, (2,)) + with self.assertRaises(StopAsyncIteration): + await g.asend(None) + + self.loop.run_until_complete(run()) + + def test_async_gen_asyncio_athrow_stopiteration(self): + async def gen(): + try: + yield 1 + except ZeroDivisionError: + yield StopIteration(2) + + async def run(): + g = gen() + v = await g.asend(None) + self.assertEqual(v, 1) + v = await g.athrow(ZeroDivisionError) + self.assertIsInstance(v, StopIteration) + self.assertEqual(v.value, 2) + with self.assertRaises(StopAsyncIteration): + await g.asend(None) + + self.loop.run_until_complete(run()) + + @unittest.skip("TODO: RUSTPYTHON") + def test_async_gen_asyncio_shutdown_01(self): + finalized = 0 + + async def waiter(timeout): + nonlocal finalized + try: + await asyncio.sleep(timeout) + yield 1 + finally: + await asyncio.sleep(0) + 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)) + + # Silence warnings + t1.cancel() + t2.cancel() + + with self.assertRaises(asyncio.CancelledError): + self.loop.run_until_complete(t1) + with self.assertRaises(asyncio.CancelledError): + self.loop.run_until_complete(t2) + + self.loop.run_until_complete(self.loop.shutdown_asyncgens()) + + self.assertEqual(finalized, 2) + + # TODO: RUSTPYTHON: async for gen expression compilation + # def test_async_gen_expression_01(self): + # async def arange(n): + # for i in range(n): + # await asyncio.sleep(0.01) + # yield i + + # def make_arange(n): + # # This syntax is legal starting with Python 3.7 + # return (i * 2 async for i in arange(n)) + + # async def run(): + # return [i async for i in make_arange(10)] + + # res = self.loop.run_until_complete(run()) + # self.assertEqual(res, [i * 2 for i in range(10)]) + + # def test_async_gen_expression_02(self): + # async def wrap(n): + # await asyncio.sleep(0.01) + # return n + + # def make_arange(n): + # # This syntax is legal starting with Python 3.7 + # return (i * 2 for i in range(n) if await wrap(i)) + + # async def run(): + # return [i async for i in make_arange(10)] + + # res = self.loop.run_until_complete(run()) + # self.assertEqual(res, [i * 2 for i in range(1, 10)]) + + def test_asyncgen_nonstarted_hooks_are_cancellable(self): + # See https://bugs.python.org/issue38013 + messages = [] + + def exception_handler(loop, context): + messages.append(context) + + async def async_iterate(): + yield 1 + yield 2 + + async def main(): + loop = asyncio.get_running_loop() + loop.set_exception_handler(exception_handler) + + async for i in async_iterate(): + break + + asyncio.run(main()) + + self.assertEqual([], messages) + + def test_async_gen_await_same_anext_coro_twice(self): + async def async_iterate(): + yield 1 + yield 2 + + async def run(): + it = async_iterate() + nxt = it.__anext__() + await nxt + with self.assertRaisesRegex( + RuntimeError, + r"cannot reuse already awaited __anext__\(\)/asend\(\)" + ): + await nxt + + await it.aclose() # prevent unfinished iterator warning + + self.loop.run_until_complete(run()) + + def test_async_gen_await_same_aclose_coro_twice(self): + async def async_iterate(): + yield 1 + yield 2 + + async def run(): + it = async_iterate() + nxt = it.aclose() + await nxt + with self.assertRaisesRegex( + RuntimeError, + r"cannot reuse already awaited aclose\(\)/athrow\(\)" + ): + await nxt + + self.loop.run_until_complete(run()) + + def test_async_gen_aclose_twice_with_different_coros(self): + # Regression test for https://bugs.python.org/issue39606 + async def async_iterate(): + yield 1 + yield 2 + + async def run(): + it = async_iterate() + await it.aclose() + await it.aclose() + + self.loop.run_until_complete(run()) + + def test_async_gen_aclose_after_exhaustion(self): + # Regression test for https://bugs.python.org/issue39606 + async def async_iterate(): + yield 1 + yield 2 + + async def run(): + it = async_iterate() + async for _ in it: + pass + await it.aclose() + + self.loop.run_until_complete(run()) + + def test_async_gen_aclose_compatible_with_get_stack(self): + async def async_generator(): + yield object() + + async def run(): + ag = async_generator() + asyncio.create_task(ag.aclose()) + tasks = asyncio.all_tasks() + for task in tasks: + # No AttributeError raised + task.get_stack() + + self.loop.run_until_complete(run()) + + +if __name__ == "__main__": + unittest.main() diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs index 034e0f9a1..8f88965b7 100644 --- a/compiler/src/compile.rs +++ b/compiler/src/compile.rs @@ -39,7 +39,7 @@ struct CompileContext { func: FunctionContext, } -#[derive(Clone, Copy)] +#[derive(Clone, Copy, PartialEq)] enum FunctionContext { NoFunction, Function, @@ -562,6 +562,14 @@ impl Compiler { } match value { Some(v) => { + if self.ctx.func == FunctionContext::AsyncFunction + && self.current_output().is_generator() + { + return Err(self.error_loc( + CompileErrorType::AsyncReturnValue, + statement.location.clone(), + )); + } self.compile_expression(v)?; } None => { @@ -1644,6 +1652,9 @@ impl Compiler { self.emit(Instruction::YieldValue); } Await { value } => { + if self.ctx.func != FunctionContext::AsyncFunction { + return Err(self.error(CompileErrorType::InvalidAwait)); + } self.compile_expression(value)?; self.emit(Instruction::GetAwaitable); self.emit(Instruction::LoadConst { @@ -1652,6 +1663,15 @@ impl Compiler { self.emit(Instruction::YieldFrom); } YieldFrom { value } => { + match self.ctx.func { + FunctionContext::NoFunction => { + return Err(self.error(CompileErrorType::InvalidYieldFrom)) + } + FunctionContext::AsyncFunction => { + return Err(self.error(CompileErrorType::AsyncYieldFrom)) + } + FunctionContext::Function => {} + } self.mark_generator(); self.compile_expression(value)?; self.emit(Instruction::GetIter); @@ -1670,7 +1690,7 @@ impl Compiler { value: bytecode::Constant::Boolean { value: false }, }); } - None => { + ast::ExpressionType::None => { self.emit(Instruction::LoadConst { value: bytecode::Constant::None, }); diff --git a/compiler/src/error.rs b/compiler/src/error.rs index 836e479ad..6f28e3a33 100644 --- a/compiler/src/error.rs +++ b/compiler/src/error.rs @@ -54,6 +54,10 @@ pub enum CompileErrorType { InvalidContinue, InvalidReturn, InvalidYield, + InvalidYieldFrom, + InvalidAwait, + AsyncYieldFrom, + AsyncReturnValue, } impl CompileError { @@ -96,6 +100,12 @@ impl fmt::Display for CompileError { CompileErrorType::InvalidContinue => "'continue' outside loop".to_owned(), CompileErrorType::InvalidReturn => "'return' outside function".to_owned(), CompileErrorType::InvalidYield => "'yield' outside function".to_owned(), + CompileErrorType::InvalidYieldFrom => "'yield from' outside function".to_owned(), + CompileErrorType::InvalidAwait => "'await' outside async function".to_owned(), + CompileErrorType::AsyncYieldFrom => "'yield from' inside async function".to_owned(), + CompileErrorType::AsyncReturnValue => { + "'return' with value inside async generator".to_owned() + } }; if let Some(statement) = &self.statement { diff --git a/compiler/src/output_stream.rs b/compiler/src/output_stream.rs index fa358817a..9cc49a944 100644 --- a/compiler/src/output_stream.rs +++ b/compiler/src/output_stream.rs @@ -7,6 +7,8 @@ pub trait OutputStream: From + Into { fn set_label(&mut self, label: Label); /// Mark the inner CodeObject as a generator fn mark_generator(&mut self); + /// Check to see if the inner CodeObject is a generator + fn is_generator(&self) -> bool; } pub struct CodeObjectStream { @@ -36,4 +38,7 @@ impl OutputStream for CodeObjectStream { fn mark_generator(&mut self) { self.code.flags |= CodeFlags::IS_GENERATOR; } + fn is_generator(&self) -> bool { + self.code.flags.contains(CodeFlags::IS_GENERATOR) + } } diff --git a/compiler/src/peephole.rs b/compiler/src/peephole.rs index b85266867..67a4b9efe 100644 --- a/compiler/src/peephole.rs +++ b/compiler/src/peephole.rs @@ -112,6 +112,9 @@ where fn mark_generator(&mut self) { self.inner.mark_generator() } + fn is_generator(&self) -> bool { + self.inner.is_generator() + } } impl OptimizationBuffer for PeepholeOptimizer { diff --git a/parser/src/lexer.rs b/parser/src/lexer.rs index 9da8825ca..f239e6e19 100644 --- a/parser/src/lexer.rs +++ b/parser/src/lexer.rs @@ -686,6 +686,8 @@ where tabs = 0; } None => { + spaces = 0; + tabs = 0; break; } _ => { diff --git a/vm/src/frame.rs b/vm/src/frame.rs index 21c54277a..922cf7cb4 100644 --- a/vm/src/frame.rs +++ b/vm/src/frame.rs @@ -7,6 +7,7 @@ use itertools::Itertools; use crate::bytecode; use crate::exceptions::{self, ExceptionCtor, PyBaseExceptionRef}; use crate::function::{single_or_tuple_any, PyFuncArgs}; +use crate::obj::objasyncgenerator::PyAsyncGenWrappedValue; use crate::obj::objbool; use crate::obj::objcode::PyCodeRef; use crate::obj::objcoroinner::Coro; @@ -118,11 +119,15 @@ impl ExecutionResult { } /// Turn an ExecutionResult into a PyResult that would be returned from a generator or coroutine - pub fn into_result(self, vm: &VirtualMachine) -> PyResult { + pub fn into_result(self, async_stopiter: bool, vm: &VirtualMachine) -> PyResult { match self { ExecutionResult::Yield(value) => Ok(value), ExecutionResult::Return(value) => { - let stop_iteration = vm.ctx.exceptions.stop_iteration.clone(); + let stop_iteration = if async_stopiter { + vm.ctx.exceptions.stop_async_iteration.clone() + } else { + vm.ctx.exceptions.stop_iteration.clone() + }; let args = if vm.is_none(&value) { vec![] } else { @@ -379,6 +384,11 @@ impl Frame { } bytecode::Instruction::YieldValue => { let value = self.pop_value(); + let value = if self.code.flags.contains(bytecode::CodeFlags::IS_COROUTINE) { + PyAsyncGenWrappedValue(value).into_ref(vm).into_object() + } else { + value + }; Ok(Some(ExecutionResult::Yield(value))) } bytecode::Instruction::YieldFrom => self.execute_yield_from(vm), diff --git a/vm/src/obj/mod.rs b/vm/src/obj/mod.rs index dd259c752..9a5523e46 100644 --- a/vm/src/obj/mod.rs +++ b/vm/src/obj/mod.rs @@ -1,5 +1,6 @@ //! This package contains the python basic/builtin types +pub mod objasyncgenerator; pub mod objbool; pub mod objbuiltinfunc; pub mod objbytearray; diff --git a/vm/src/obj/objasyncgenerator.rs b/vm/src/obj/objasyncgenerator.rs new file mode 100644 index 000000000..e209040d0 --- /dev/null +++ b/vm/src/obj/objasyncgenerator.rs @@ -0,0 +1,401 @@ +use super::objcode::PyCodeRef; +use super::objcoroinner::{Coro, Variant}; +use super::objtype::{self, PyClassRef}; +use crate::exceptions::PyBaseExceptionRef; +use crate::frame::FrameRef; +use crate::function::OptionalArg; +use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; +use crate::vm::VirtualMachine; + +use std::cell::Cell; + +#[pyclass(name = "async_generator")] +#[derive(Debug)] +pub struct PyAsyncGen { + inner: Coro, + running_async: Cell, +} +pub type PyAsyncGenRef = PyRef; + +impl PyValue for PyAsyncGen { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.ctx.types.async_generator.clone() + } +} + +#[pyimpl] +impl PyAsyncGen { + pub fn as_coro(&self) -> &Coro { + &self.inner + } + + pub fn new(frame: FrameRef, vm: &VirtualMachine) -> PyAsyncGenRef { + PyAsyncGen { + inner: Coro::new(frame, Variant::AsyncGen), + running_async: Cell::new(false), + } + .into_ref(vm) + } + + // TODO: fix function names situation + #[pyproperty(magic)] + fn name(&self, vm: &VirtualMachine) -> PyObjectRef { + vm.get_none() + } + + #[pymethod(name = "__aiter__")] + fn aiter(zelf: PyRef, _vm: &VirtualMachine) -> PyRef { + zelf + } + + #[pymethod(name = "__anext__")] + fn anext(zelf: PyRef, vm: &VirtualMachine) -> PyAsyncGenASend { + Self::asend(zelf, vm.get_none(), vm) + } + + #[pymethod] + fn asend(zelf: PyRef, value: PyObjectRef, _vm: &VirtualMachine) -> PyAsyncGenASend { + PyAsyncGenASend { + ag: zelf, + state: Cell::new(AwaitableState::Init), + value, + } + } + + #[pymethod] + fn athrow( + zelf: PyRef, + exc_type: PyObjectRef, + exc_val: OptionalArg, + exc_tb: OptionalArg, + vm: &VirtualMachine, + ) -> PyAsyncGenAThrow { + PyAsyncGenAThrow { + ag: zelf, + aclose: false, + state: Cell::new(AwaitableState::Init), + value: ( + exc_type, + exc_val.unwrap_or_else(|| vm.get_none()), + exc_tb.unwrap_or_else(|| vm.get_none()), + ), + } + } + + #[pymethod] + fn aclose(zelf: PyRef, vm: &VirtualMachine) -> PyAsyncGenAThrow { + PyAsyncGenAThrow { + ag: zelf, + aclose: true, + state: Cell::new(AwaitableState::Init), + value: ( + vm.ctx.exceptions.generator_exit.clone().into_object(), + vm.get_none(), + vm.get_none(), + ), + } + } + + #[pyproperty] + fn ag_await(&self, _vm: &VirtualMachine) -> Option { + self.inner.frame().yield_from_target() + } + #[pyproperty] + fn ag_frame(&self, _vm: &VirtualMachine) -> FrameRef { + self.inner.frame() + } + #[pyproperty] + fn ag_running(&self, _vm: &VirtualMachine) -> bool { + self.inner.running() + } + #[pyproperty] + fn ag_code(&self, _vm: &VirtualMachine) -> PyCodeRef { + self.inner.frame().code.clone() + } +} + +#[pyclass(name = "async_generator_wrapped_value")] +#[derive(Debug)] +pub(crate) struct PyAsyncGenWrappedValue(pub PyObjectRef); +impl PyValue for PyAsyncGenWrappedValue { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.ctx.types.async_generator_wrapped_value.clone() + } +} + +impl PyAsyncGenWrappedValue { + fn unbox(ag: &PyAsyncGen, val: PyResult, vm: &VirtualMachine) -> PyResult { + if let Err(ref e) = val { + if objtype::isinstance(&e, &vm.ctx.exceptions.stop_async_iteration) + || objtype::isinstance(&e, &vm.ctx.exceptions.generator_exit) + { + ag.inner.closed.set(true); + } + ag.running_async.set(false); + } + let val = val?; + + match_class!(match val { + val @ Self => { + ag.running_async.set(false); + Err(vm.new_exception( + vm.ctx.exceptions.stop_iteration.clone(), + vec![val.0.clone()], + )) + } + val => Ok(val), + }) + } +} + +#[derive(Debug, Clone, Copy)] +enum AwaitableState { + Init, + Iter, + Closed, +} + +#[pyclass(name = "async_generator_asend")] +#[derive(Debug)] +struct PyAsyncGenASend { + ag: PyAsyncGenRef, + state: Cell, + value: PyObjectRef, +} + +impl PyValue for PyAsyncGenASend { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.ctx.types.async_generator_asend.clone() + } +} + +#[pyimpl] +impl PyAsyncGenASend { + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef, _vm: &VirtualMachine) -> PyRef { + zelf + } + #[pymethod(name = "__await__")] + fn r#await(zelf: PyRef, _vm: &VirtualMachine) -> PyRef { + zelf + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + self.send(vm.get_none(), vm) + } + + #[pymethod] + fn send(&self, val: PyObjectRef, vm: &VirtualMachine) -> PyResult { + let val = match self.state.get() { + AwaitableState::Closed => { + return Err(vm.new_runtime_error( + "cannot reuse already awaited __anext__()/asend()".to_owned(), + )) + } + AwaitableState::Iter => val, // already running, all good + AwaitableState::Init => { + if self.ag.running_async.get() { + return Err(vm.new_runtime_error( + "anext(): asynchronous generator is already running".to_owned(), + )); + } + self.ag.running_async.set(true); + self.state.set(AwaitableState::Iter); + if vm.is_none(&val) { + self.value.clone() + } else { + val + } + } + }; + let res = self.ag.inner.send(val, vm); + let res = PyAsyncGenWrappedValue::unbox(&self.ag, res, vm); + if res.is_err() { + self.close(); + } + res + } + + #[pymethod] + fn throw( + &self, + exc_type: PyObjectRef, + exc_val: OptionalArg, + exc_tb: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult { + if let AwaitableState::Closed = self.state.get() { + return Err( + vm.new_runtime_error("cannot reuse already awaited __anext__()/asend()".to_owned()) + ); + } + + let res = self.ag.inner.throw( + exc_type, + exc_val.unwrap_or_else(|| vm.get_none()), + exc_tb.unwrap_or_else(|| vm.get_none()), + vm, + ); + let res = PyAsyncGenWrappedValue::unbox(&self.ag, res, vm); + if res.is_err() { + self.close(); + } + res + } + + #[pymethod] + fn close(&self) { + self.state.set(AwaitableState::Closed); + } +} + +#[pyclass(name = "async_generator_athrow")] +#[derive(Debug)] +struct PyAsyncGenAThrow { + ag: PyAsyncGenRef, + aclose: bool, + state: Cell, + value: (PyObjectRef, PyObjectRef, PyObjectRef), +} + +impl PyValue for PyAsyncGenAThrow { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.ctx.types.async_generator_athrow.clone() + } +} + +#[pyimpl] +impl PyAsyncGenAThrow { + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef, _vm: &VirtualMachine) -> PyRef { + zelf + } + #[pymethod(name = "__await__")] + fn r#await(zelf: PyRef, _vm: &VirtualMachine) -> PyRef { + zelf + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + self.send(vm.get_none(), vm) + } + + #[pymethod] + fn send(&self, val: PyObjectRef, vm: &VirtualMachine) -> PyResult { + match self.state.get() { + AwaitableState::Closed => { + Err(vm + .new_runtime_error("cannot reuse already awaited aclose()/athrow()".to_owned())) + } + AwaitableState::Init => { + if self.ag.running_async.get() { + self.state.set(AwaitableState::Closed); + let msg = if self.aclose { + "aclose(): asynchronous generator is already running" + } else { + "athrow(): asynchronous generator is already running" + }; + return Err(vm.new_runtime_error(msg.to_owned())); + } + if self.ag.inner.closed() { + self.state.set(AwaitableState::Closed); + return Err(vm.new_exception_empty(vm.ctx.exceptions.stop_iteration.clone())); + } + if !vm.is_none(&val) { + return Err(vm.new_runtime_error( + "can't send non-None value to a just-started async generator".to_owned(), + )); + } + self.state.set(AwaitableState::Iter); + self.ag.running_async.set(true); + + let (ty, val, tb) = self.value.clone(); + let ret = self.ag.inner.throw(ty, val, tb, vm); + let ret = if self.aclose { + if self.ignored_close(&ret) { + Err(self.yield_close(vm)) + } else { + ret + } + } else { + PyAsyncGenWrappedValue::unbox(&self.ag, ret, vm) + }; + ret.map_err(|e| self.check_error(e, vm)) + } + AwaitableState::Iter => { + let ret = self.ag.inner.send(val, vm); + if self.aclose { + match ret { + Ok(v) if v.payload_is::() => { + Err(self.yield_close(vm)) + } + Ok(v) => Ok(v), + Err(e) => Err(self.check_error(e, vm)), + } + } else { + PyAsyncGenWrappedValue::unbox(&self.ag, ret, vm) + } + } + } + } + + #[pymethod] + fn throw( + &self, + exc_type: PyObjectRef, + exc_val: OptionalArg, + exc_tb: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult { + let ret = self.ag.inner.throw( + exc_type, + exc_val.unwrap_or_else(|| vm.get_none()), + exc_tb.unwrap_or_else(|| vm.get_none()), + vm, + ); + let res = if self.aclose { + if self.ignored_close(&ret) { + Err(self.yield_close(vm)) + } else { + ret + } + } else { + PyAsyncGenWrappedValue::unbox(&self.ag, ret, vm) + }; + res.map_err(|e| self.check_error(e, vm)) + } + + #[pymethod] + fn close(&self) { + self.state.set(AwaitableState::Closed); + } + + fn ignored_close(&self, res: &PyResult) -> bool { + res.as_ref() + .map_or(false, |v| v.payload_is::()) + } + fn yield_close(&self, vm: &VirtualMachine) -> PyBaseExceptionRef { + self.ag.running_async.set(false); + self.state.set(AwaitableState::Closed); + vm.new_runtime_error("async generator ignored GeneratorExit".to_owned()) + } + fn check_error(&self, exc: PyBaseExceptionRef, vm: &VirtualMachine) -> PyBaseExceptionRef { + self.ag.running_async.set(false); + self.state.set(AwaitableState::Closed); + if self.aclose + && (objtype::isinstance(&exc, &vm.ctx.exceptions.stop_async_iteration) + || objtype::isinstance(&exc, &vm.ctx.exceptions.generator_exit)) + { + vm.new_exception_empty(vm.ctx.exceptions.stop_iteration.clone()) + } else { + exc + } + } +} + +pub fn init(ctx: &PyContext) { + PyAsyncGen::extend_class(ctx, &ctx.types.async_generator); + PyAsyncGenASend::extend_class(ctx, &ctx.types.async_generator_asend); + PyAsyncGenAThrow::extend_class(ctx, &ctx.types.async_generator_athrow); +} diff --git a/vm/src/obj/objcoroinner.rs b/vm/src/obj/objcoroinner.rs index dfb2ab264..97cd810e8 100644 --- a/vm/src/obj/objcoroinner.rs +++ b/vm/src/obj/objcoroinner.rs @@ -1,4 +1,4 @@ -use super::{objiter, objtype}; +use super::objtype::{self, PyClassRef}; use crate::exceptions::{self, PyBaseExceptionRef}; use crate::frame::{ExecutionResult, FrameRef}; use crate::pyobject::{PyObjectRef, PyResult}; @@ -6,21 +6,50 @@ use crate::vm::VirtualMachine; use std::cell::{Cell, RefCell}; +#[derive(Debug, PartialEq, Clone, Copy)] +pub enum Variant { + Gen, + Coroutine, + AsyncGen, +} +impl Variant { + fn exec_result(self, res: ExecutionResult, vm: &VirtualMachine) -> PyResult { + res.into_result(self == Self::AsyncGen, vm) + } + fn name(self) -> &'static str { + match self { + Self::Gen => "generator", + Self::Coroutine => "coroutine", + Self::AsyncGen => "async generator", + } + } + fn stop_iteration(self, vm: &VirtualMachine) -> PyClassRef { + match self { + Self::AsyncGen => vm.ctx.exceptions.stop_async_iteration.clone(), + _ => vm.ctx.exceptions.stop_iteration.clone(), + } + } +} + #[derive(Debug)] pub struct Coro { frame: FrameRef, - closed: Cell, + pub closed: Cell, running: Cell, exceptions: RefCell>, + started: Cell, + variant: Variant, } impl Coro { - pub fn new(frame: FrameRef) -> Self { + pub fn new(frame: FrameRef, variant: Variant) -> Self { Coro { frame, closed: Cell::new(false), running: Cell::new(false), exceptions: RefCell::new(vec![]), + started: Cell::new(false), + variant, } } @@ -47,20 +76,44 @@ impl Coro { .split_off(curr_exception_stack_len), ); self.running.set(false); + self.started.set(true); result } pub fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult { if self.closed.get() { - return Err(objiter::new_stop_iteration(vm)); + return Err(vm.new_exception_empty(self.variant.stop_iteration(vm))); + } + if !self.started.get() && !vm.is_none(&value) { + return Err(vm.new_type_error(format!( + "can't send non-None value to a just-started {}", + self.variant.name() + ))); } - self.frame.push_value(value.clone()); let result = self.run_with_context(|| vm.run_frame(self.frame.clone()), vm); self.maybe_close(&result); - result?.into_result(vm) + match result { + Ok(exec_res) => self.variant.exec_result(exec_res, vm), + Err(e) => { + if objtype::isinstance(&e, &vm.ctx.exceptions.stop_iteration) { + let err = vm + .new_runtime_error(format!("{} raised StopIteration", self.variant.name())); + err.set_cause(Some(e)); + Err(err) + } else if self.variant == Variant::AsyncGen + && objtype::isinstance(&e, &vm.ctx.exceptions.stop_async_iteration) + { + let err = vm + .new_runtime_error("async generator raised StopAsyncIteration".to_owned()); + err.set_cause(Some(e)); + Err(err) + } else { + Err(e) + } + } + } } - pub fn throw( &self, exc_type: PyObjectRef, @@ -76,7 +129,7 @@ impl Coro { self.run_with_context(|| self.frame.gen_throw(vm, exc_type, exc_val, exc_tb), vm); self.maybe_close(&result); vm.frames.borrow_mut().pop(); - result?.into_result(vm) + self.variant.exec_result(result?, vm) } pub fn close(&self, vm: &VirtualMachine) -> PyResult<()> { @@ -99,22 +152,28 @@ impl Coro { self.closed.set(true); match result { Ok(ExecutionResult::Yield(_)) => { - Err(vm.new_runtime_error("generator ignored GeneratorExit".to_owned())) + Err(vm.new_runtime_error(format!("{} ignored GeneratorExit", self.variant.name()))) } Err(e) if !is_gen_exit(&e, vm) => Err(e), _ => Ok(()), } } - pub fn closed(&self) -> bool { - self.closed.get() + pub fn started(&self) -> bool { + self.started.get() } pub fn running(&self) -> bool { self.running.get() } + pub fn closed(&self) -> bool { + self.closed.get() + } pub fn frame(&self) -> FrameRef { self.frame.clone() } + pub fn name(&self) -> String { + self.frame.code.obj_name.clone() + } } pub fn is_gen_exit(exc: &PyBaseExceptionRef, vm: &VirtualMachine) -> bool { diff --git a/vm/src/obj/objcoroutine.rs b/vm/src/obj/objcoroutine.rs index ac7a77fb2..d1faacc2b 100644 --- a/vm/src/obj/objcoroutine.rs +++ b/vm/src/obj/objcoroutine.rs @@ -1,5 +1,5 @@ use super::objcode::PyCodeRef; -use super::objcoroinner::Coro; +use super::objcoroinner::{Coro, Variant}; use super::objstr::PyStringRef; use super::objtype::PyClassRef; use crate::frame::FrameRef; @@ -29,11 +29,17 @@ impl PyCoroutine { pub fn new(frame: FrameRef, vm: &VirtualMachine) -> PyCoroutineRef { PyCoroutine { - inner: Coro::new(frame), + inner: Coro::new(frame, Variant::Coroutine), } .into_ref(vm) } + // TODO: fix function names situation + #[pyproperty(magic)] + fn name(&self, vm: &VirtualMachine) -> PyObjectRef { + vm.get_none() + } + #[pymethod] fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult { self.inner.send(value, vm) diff --git a/vm/src/obj/objfunction.rs b/vm/src/obj/objfunction.rs index b9994712d..bb3af08f4 100644 --- a/vm/src/obj/objfunction.rs +++ b/vm/src/obj/objfunction.rs @@ -6,6 +6,7 @@ use super::objtype::PyClassRef; use crate::bytecode; use crate::frame::Frame; use crate::function::{OptionalArg, PyFuncArgs}; +use crate::obj::objasyncgenerator::PyAsyncGen; use crate::obj::objcoroutine::PyCoroutine; use crate::obj::objgenerator::PyGenerator; use crate::pyobject::{ @@ -221,12 +222,13 @@ impl PyFunction { let frame = Frame::new(code.clone(), scope).into_ref(vm); // If we have a generator, create a new generator - if code.flags.contains(bytecode::CodeFlags::IS_GENERATOR) { - Ok(PyGenerator::new(frame, vm).into_object()) - } else if code.flags.contains(bytecode::CodeFlags::IS_COROUTINE) { - Ok(PyCoroutine::new(frame, vm).into_object()) - } else { - vm.run_frame_full(frame) + let is_gen = code.flags.contains(bytecode::CodeFlags::IS_GENERATOR); + let is_coro = code.flags.contains(bytecode::CodeFlags::IS_COROUTINE); + match (is_gen, is_coro) { + (true, false) => Ok(PyGenerator::new(frame, vm).into_object()), + (false, true) => Ok(PyCoroutine::new(frame, vm).into_object()), + (true, true) => Ok(PyAsyncGen::new(frame, vm).into_object()), + (false, false) => vm.run_frame_full(frame), } } diff --git a/vm/src/obj/objgenerator.rs b/vm/src/obj/objgenerator.rs index 45b81de99..e91e24406 100644 --- a/vm/src/obj/objgenerator.rs +++ b/vm/src/obj/objgenerator.rs @@ -3,7 +3,7 @@ */ use super::objcode::PyCodeRef; -use super::objcoroinner::Coro; +use super::objcoroinner::{Coro, Variant}; use super::objtype::PyClassRef; use crate::frame::FrameRef; use crate::function::OptionalArg; @@ -32,11 +32,17 @@ impl PyGenerator { pub fn new(frame: FrameRef, vm: &VirtualMachine) -> PyGeneratorRef { PyGenerator { - inner: Coro::new(frame), + inner: Coro::new(frame, Variant::Gen), } .into_ref(vm) } + // TODO: fix function names situation + #[pyproperty(magic)] + fn name(&self, vm: &VirtualMachine) -> PyObjectRef { + vm.get_none() + } + #[pymethod(name = "__iter__")] fn iter(zelf: PyGeneratorRef) -> PyGeneratorRef { zelf diff --git a/vm/src/types.rs b/vm/src/types.rs index fc8caec05..9e3b6a469 100644 --- a/vm/src/types.rs +++ b/vm/src/types.rs @@ -1,3 +1,4 @@ +use crate::obj::objasyncgenerator; use crate::obj::objbool; use crate::obj::objbuiltinfunc; use crate::obj::objbytearray; @@ -47,6 +48,10 @@ use std::rc::Rc; /// Holder of references to builtin types. #[derive(Debug)] pub struct TypeZoo { + pub async_generator: PyClassRef, + pub async_generator_asend: PyClassRef, + pub async_generator_athrow: PyClassRef, + pub async_generator_wrapped_value: PyClassRef, pub bytes_type: PyClassRef, pub bytesiterator_type: PyClassRef, pub bytearray_type: PyClassRef, @@ -134,6 +139,12 @@ impl TypeZoo { let generator_type = create_type("generator", &type_type, &object_type); let coroutine_type = create_type("coroutine", &type_type, &object_type); let coroutine_wrapper_type = create_type("coroutine_wrapper", &type_type, &object_type); + let async_generator = create_type("async_generator", &type_type, &object_type); + let async_generator_asend = create_type("async_generator_asend", &type_type, &object_type); + let async_generator_athrow = + create_type("async_generator_athrow", &type_type, &object_type); + let async_generator_wrapped_value = + create_type("async_generator_wrapped_value", &type_type, &object_type); let bound_method_type = create_type("method", &type_type, &object_type); let str_type = create_type("str", &type_type, &object_type); let list_type = create_type("list", &type_type, &object_type); @@ -175,6 +186,10 @@ impl TypeZoo { let traceback_type = create_type("traceback", &type_type, &object_type); Self { + async_generator, + async_generator_asend, + async_generator_athrow, + async_generator_wrapped_value, bool_type, memoryview_type, bytearray_type, @@ -379,6 +394,7 @@ pub fn initialize_types(context: &PyContext) { objclassmethod::init(&context); objgenerator::init(&context); objcoroutine::init(&context); + objasyncgenerator::init(&context); objint::init(&context); objfloat::init(&context); objcomplex::init(&context);