Address review comments.

cjerdonek · cjerdonek · commit 66c71aa01881 · 2020-05-21T20:36:47.000-07:00
diff --git a/Include/internal/pycore_pyerrors.h b/Include/internal/pycore_pyerrors.h
@@ -50,8 +50,8 @@ PyAPI_FUNC(void) _PyErr_SetObject(
     PyObject *type,
     PyObject *value);
 
-PyAPI_FUNC(void) _PyErr_ChainThreadState(
-    PyThreadState *tstate);
+PyAPI_FUNC(void) _PyErr_ChainStackItem(
+    _PyErr_StackItem *exc_info);
 
 PyAPI_FUNC(void) _PyErr_Clear(PyThreadState *tstate);
 
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
@@ -536,9 +536,11 @@ async def run():
                     self.assertEqual((type(chained), chained.args),
                         (KeyError, (3,)))
 
-        task = self.new_task(loop, run())
-        loop.run_until_complete(task)
-        loop.close()
+        try:
+            task = self.new_task(loop, run())
+            loop.run_until_complete(task)
+        finally:
+            loop.close()
 
     def test_exception_chaining_after_await_with_context_cycle(self):
         # Check trying to create an exception context cycle:
@@ -563,9 +565,12 @@ async def run():
                     # Prevent a hang if has_cycle is True.
                     exc.__context__ = None
 
-        task = self.new_task(loop, run())
-        loop.run_until_complete(task)
-        loop.close()
+        try:
+            task = self.new_task(loop, run())
+            loop.run_until_complete(task)
+        finally:
+            loop.close()
+        # This also distinguishes from the initial has_cycle=None.
         self.assertEqual(has_cycle, False)
 
     def test_cancel(self):
diff --git a/Lib/test/test_generators.py b/Lib/test/test_generators.py
@@ -394,6 +394,7 @@ def g(exc):
         gen = g(exc)
         gen.send(None)
         gen.throw(exc)
+        # This also distinguishes from the initial has_cycle=None.
         self.assertEqual(has_cycle, False)
 
     def test_throw_after_none_exc_type(self):
diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
@@ -1,6 +1,5 @@
 #include "Python.h"
 #include "pycore_pyerrors.h"      // _PyErr_ClearExcState()
-#include "pycore_pystate.h"       // _PyThreadState_GET()
 #include <stddef.h>               // offsetof()
 
 
@@ -619,12 +618,7 @@ future_set_cancelled_error(FutureObj *fut)
     PyErr_SetObject(asyncio_CancelledError, exc);
     Py_DECREF(exc);
 
-    PyThreadState *tstate = _PyThreadState_GET();
-
-    _PyErr_StackItem *saved_exc_info = tstate->exc_info;
-    tstate->exc_info = &fut->fut_cancelled_exc_state;
-    _PyErr_ChainThreadState(tstate);
-    tstate->exc_info = saved_exc_info;
+    _PyErr_ChainStackItem(&fut->fut_cancelled_exc_state);
 }
 
 static int
diff --git a/Objects/genobject.c b/Objects/genobject.c
@@ -207,9 +207,9 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing)
     gen->gi_exc_state.previous_item = tstate->exc_info;
     tstate->exc_info = &gen->gi_exc_state;
 
-    /* XXX: can we assert that _PyErr_Occurred() is true when exc is true? */
-    if (exc && _PyErr_Occurred(tstate)) {
-        _PyErr_ChainThreadState(tstate);
+    if (exc) {
+        assert(_PyErr_Occurred(tstate));
+        _PyErr_ChainStackItem(NULL);
     }
 
     result = _PyEval_EvalFrame(tstate, f, exc);
diff --git a/Python/errors.c b/Python/errors.c
@@ -514,39 +514,60 @@ _PyErr_ChainExceptions(PyObject *exc, PyObject *val, PyObject *tb)
     }
 }
 
-/* Set the context of the currently set exception to PyThreadState's
-   topmost exception.
+/* Set the currently set exception's context to the given exception.
 
-   This function won't create any cycles in the exception context chain
-   to the extent that _PyErr_SetObject ensures this.
+   If the provided exc_info is NULL, then the current Python thread state's
+   exc_info will be used for the context instead.
 
-   This function can only be called when _PyErr_Occurred() is true. */
+   This function can only be called when _PyErr_Occurred() is true.
+   Also, this function won't create any cycles in the exception context
+   chain to the extent that _PyErr_SetObject ensures this. */
 void
-_PyErr_ChainThreadState(PyThreadState *tstate)
+_PyErr_ChainStackItem(_PyErr_StackItem *exc_info)
 {
+    PyThreadState *tstate = _PyThreadState_GET();
     assert(_PyErr_Occurred(tstate));
-    _PyErr_StackItem *exc_info = tstate->exc_info;
 
+    int exc_info_given;
+    if (exc_info == NULL) {
+        exc_info_given = 0;
+        exc_info = tstate->exc_info;
+    } else {
+        exc_info_given = 1;
+    }
     if (exc_info->exc_type == NULL || exc_info->exc_type == Py_None) {
         return;
     }
 
-    PyObject *exc, *val, *tb;
-    exc = exc_info->exc_type;
-    val = exc_info->exc_value;
-    tb = exc_info->exc_traceback;
-    _PyErr_NormalizeException(tstate, &exc, &val, &tb);
-    if (tb != NULL) {
-        PyException_SetTraceback(val, tb);
+    _PyErr_StackItem *saved_exc_info;
+    if (exc_info_given) {
+        /* Temporarily set the thread state's exc_info since this is what
+           _PyErr_SetObject uses for implicit exception chaining. */
+        saved_exc_info = tstate->exc_info;
+        tstate->exc_info = exc_info;
     }
 
+    PyObject *exc, *val, *tb;
+    _PyErr_Fetch(tstate, &exc, &val, &tb);
+
     PyObject *exc2, *val2, *tb2;
-    _PyErr_Fetch(tstate, &exc2, &val2, &tb2);
+    exc2 = exc_info->exc_type;
+    val2 = exc_info->exc_value;
+    tb2 = exc_info->exc_traceback;
+    _PyErr_NormalizeException(tstate, &exc2, &val2, &tb2);
+    if (tb2 != NULL) {
+        PyException_SetTraceback(val2, tb2);
+    }
+
     /* _PyErr_SetObject sets the context from PyThreadState. */
-    _PyErr_SetObject(tstate, exc2, val2);
-    Py_XDECREF(exc2);
-    Py_XDECREF(val2);
-    Py_XDECREF(tb2);
+    _PyErr_SetObject(tstate, exc, val);
+    Py_DECREF(exc);  // since _PyErr_Occurred was true
+    Py_XDECREF(val);
+    Py_XDECREF(tb);
+
+    if (exc_info_given) {
+        tstate->exc_info = saved_exc_info;
+    }
 }
 
 static PyObject *
