diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py index 68f3b8cce9f65d..6eb6b46ec8af75 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/Lib/test/test_asyncio/test_tasks.py @@ -466,6 +466,33 @@ async def inner2(): t = outer() self.assertEqual(self.loop.run_until_complete(t), 1042) + def test_exception_chaining_after_await(self): + # Test that when awaiting on a task when an exception is already + # active, if the task raises an exception it will be chained + # with the original. + loop = asyncio.new_event_loop() + self.set_event_loop(loop) + + async def raise_error(): + raise ValueError + + async def run(): + try: + raise KeyError(3) + except Exception as exc: + task = self.new_task(loop, raise_error()) + try: + await task + except Exception as exc: + self.assertEqual(type(exc), ValueError) + chained = exc.__context__ + self.assertEqual((type(chained), chained.args), + (KeyError, (3,))) + + task = self.new_task(loop, run()) + loop.run_until_complete(task) + loop.close() + def test_cancel(self): def gen(): diff --git a/Lib/test/test_generators.py b/Lib/test/test_generators.py index e0478011996807..1081107ee64ace 100644 --- a/Lib/test/test_generators.py +++ b/Lib/test/test_generators.py @@ -318,7 +318,7 @@ def g(): class GeneratorThrowTest(unittest.TestCase): - def test_exception_context_set(self): + def test_exception_context_with_yield(self): def f(): try: raise KeyError('a') @@ -332,6 +332,23 @@ def f(): context = cm.exception.__context__ self.assertEqual((type(context), context.args), (KeyError, ('a',))) + def test_exception_context_with_yield_from(self): + def f(): + yield + + def g(): + try: + raise KeyError('a') + except Exception: + yield from f() + + gen = g() + gen.send(None) + with self.assertRaises(ValueError) as cm: + gen.throw(ValueError) + context = cm.exception.__context__ + self.assertEqual((type(context), context.args), (KeyError, ('a',))) + def test_throw_after_none_exc_type(self): def g(): try: diff --git a/Objects/genobject.c b/Objects/genobject.c index 5b253edfdcd0f6..fb01e581f8ae15 100644 --- a/Objects/genobject.c +++ b/Objects/genobject.c @@ -217,6 +217,18 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing) assert(f->f_back == NULL); f->f_back = tstate->frame; + _PyErr_StackItem *gi_exc_state = &gen->gi_exc_state; + if (exc && gi_exc_state->exc_type != NULL && + gi_exc_state->exc_type != Py_None) + { + Py_INCREF(gi_exc_state->exc_type); + Py_XINCREF(gi_exc_state->exc_value); + Py_XINCREF(gi_exc_state->exc_traceback); + _PyErr_ChainExceptions(gi_exc_state->exc_type, + gi_exc_state->exc_value, + gi_exc_state->exc_traceback); + } + gen->gi_running = 1; gen->gi_exc_state.previous_item = tstate->exc_info; tstate->exc_info = &gen->gi_exc_state; @@ -512,16 +524,6 @@ _gen_throw(PyGenObject *gen, int close_on_genexit, } PyErr_Restore(typ, val, tb); - - _PyErr_StackItem *gi_exc_state = &gen->gi_exc_state; - if (gi_exc_state->exc_type != NULL && gi_exc_state->exc_type != Py_None) { - Py_INCREF(gi_exc_state->exc_type); - Py_XINCREF(gi_exc_state->exc_value); - Py_XINCREF(gi_exc_state->exc_traceback); - _PyErr_ChainExceptions(gi_exc_state->exc_type, - gi_exc_state->exc_value, - gi_exc_state->exc_traceback); - } return gen_send_ex(gen, Py_None, 1, 0); failed_throw: