bpo-45924: Fix asyncio incorrect traceback when future's exception is raised multiple times

kumaraditya303 · kumaraditya303 · commit 15a9cfebdf10 · 2022-01-24T12:39:00.000+05:30
diff --git a/Lib/asyncio/futures.py b/Lib/asyncio/futures.py
@@ -198,7 +198,7 @@ def result(self):
             raise exceptions.InvalidStateError('Result is not ready.')
         self.__log_traceback = False
         if self._exception is not None:
-            raise self._exception
+            raise self._exception.with_traceback(self._exception_tb)
         return self._result
 
     def exception(self):
@@ -274,6 +274,7 @@ def set_exception(self, exception):
             raise TypeError("StopIteration interacts badly with generators "
                             "and cannot be raised into a Future")
         self._exception = exception
+        self._exception_tb = exception.__traceback__
         self._state = _FINISHED
         self.__schedule_callbacks()
         self.__log_traceback = True
diff --git a/Lib/test/test_asyncio/test_futures2.py b/Lib/test/test_asyncio/test_futures2.py
@@ -1,9 +1,32 @@
 # IsolatedAsyncioTestCase based tests
 import asyncio
 import unittest
-
+import traceback
 
 class FutureTests(unittest.IsolatedAsyncioTestCase):
+    async def test_future_traceback(self): 
+        async def raise_exc():
+            raise TypeError(42)
+        future = asyncio.create_task(raise_exc())
+        for _ in range(10):
+            try:
+                await future
+            except TypeError:
+                tb = traceback.format_exc()
+                expected = (
+f"""Traceback (most recent call last):
+  File "{__file__}", line 13, in test_future_traceback
+    await future
+    ^^^^^^^^^^^^
+  File "{__file__}", line 9, in raise_exc
+    raise TypeError(42)
+    ^^^^^^^^^^^^^^^^^^^
+TypeError: 42
+"""
+        )
+            self.assertEqual(tb, expected)
+
+        
     async def test_recursive_repr_for_pending_tasks(self):
         # The call crashes if the guard for recursive call
         # in base_futures:_future_repr_info is absent
diff --git a/Misc/NEWS.d/next/Library/2021-12-27-15-32-15.bpo-45924.0ZpHX2.rst b/Misc/NEWS.d/next/Library/2021-12-27-15-32-15.bpo-45924.0ZpHX2.rst
@@ -0,0 +1 @@
+Fix :mod:`asyncio` incorrect traceback when future's exception is raised multiple times. Patch by Kumar Aditya.
diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
@@ -68,6 +68,7 @@ typedef enum {
     PyObject *prefix##_context0;                                            \
     PyObject *prefix##_callbacks;                                           \
     PyObject *prefix##_exception;                                           \
+    PyObject *prefix##_exception_tb;                                        \
     PyObject *prefix##_result;                                              \
     PyObject *prefix##_source_tb;                                           \
     PyObject *prefix##_cancel_msg;                                          \
@@ -492,6 +493,7 @@ future_init(FutureObj *fut, PyObject *loop)
     Py_CLEAR(fut->fut_callbacks);
     Py_CLEAR(fut->fut_result);
     Py_CLEAR(fut->fut_exception);
+    Py_CLEAR(fut->fut_exception_tb);
     Py_CLEAR(fut->fut_source_tb);
     Py_CLEAR(fut->fut_cancel_msg);
     _PyErr_ClearExcState(&fut->fut_cancelled_exc_state);
@@ -598,7 +600,9 @@ future_set_exception(FutureObj *fut, PyObject *exc)
     }
 
     assert(!fut->fut_exception);
+    assert(!fut->fut_exception_tb);
     fut->fut_exception = exc_val;
+    fut->fut_exception_tb = PyException_GetTraceback(exc_val);
     fut->fut_state = STATE_FINISHED;
 
     if (future_schedule_callbacks(fut) == -1) {
@@ -646,6 +650,15 @@ future_get_result(FutureObj *fut, PyObject **result)
 
     fut->fut_log_tb = 0;
     if (fut->fut_exception != NULL) {
+        if (fut->fut_exception_tb != NULL) {
+            if (PyException_SetTraceback(fut->fut_exception, fut->fut_exception_tb) < 0)
+            {
+                return -1;
+            }
+        }
+        else {
+            PyException_SetTraceback(fut->fut_exception, Py_None);
+        }
         Py_INCREF(fut->fut_exception);
         *result = fut->fut_exception;
         return 1;
@@ -789,6 +802,7 @@ FutureObj_clear(FutureObj *fut)
     Py_CLEAR(fut->fut_callbacks);
     Py_CLEAR(fut->fut_result);
     Py_CLEAR(fut->fut_exception);
+    Py_CLEAR(fut->fut_exception_tb);
     Py_CLEAR(fut->fut_source_tb);
     Py_CLEAR(fut->fut_cancel_msg);
     _PyErr_ClearExcState(&fut->fut_cancelled_exc_state);
@@ -805,6 +819,7 @@ FutureObj_traverse(FutureObj *fut, visitproc visit, void *arg)
     Py_VISIT(fut->fut_callbacks);
     Py_VISIT(fut->fut_result);
     Py_VISIT(fut->fut_exception);
+    Py_VISIT(fut->fut_exception_tb);
     Py_VISIT(fut->fut_source_tb);
     Py_VISIT(fut->fut_cancel_msg);
     Py_VISIT(fut->dict);

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+Fix :mod:`asyncio` incorrect traceback when future's exception is raised multiple times. Patch by Kumar Aditya.