From 75cd8e48c62c97fdb9d9a94fd2335be06084471d Mon Sep 17 00:00:00 2001
From: Chris Jerdonek <chris.jerdonek@gmail.com>
Date: Wed, 13 May 2020 16:18:27 -0700
Subject: [PATCH] bpo-29587: Make gen.throw() chain exceptions with yield from
 (GH-19858)

The previous commits on bpo-29587 got exception chaining working
with gen.throw() in the `yield` case. This patch also gets the
`yield from` case working.

As a consequence, implicit exception chaining now also works in
the asyncio scenario of awaiting on a task when an exception is
already active.

Tests are included for both the asyncio case and the pure
generator-only case.
---
 Lib/test/test_asyncio/test_tasks.py | 27 +++++++++++++++++++++++++++
 Lib/test/test_generators.py         | 19 ++++++++++++++++++-
 Objects/genobject.c                 | 22 ++++++++++++----------
 3 files changed, 57 insertions(+), 11 deletions(-)

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: