bpo-31033: Improve the traceback for cancelled asyncio tasks (GH-19951)

When an asyncio.Task is cancelled, the exception traceback now
starts with where the task was first interrupted.  Previously,
the traceback only had "depth one."

1 files changed, 21 insertions, 5 deletions

diff --git a/Lib/asyncio/futures.py b/Lib/asyncio/futures.py
index 889f3e6..bed4da5 100644
--- a/Lib/asyncio/futures.py
+++ b/Lib/asyncio/futures.py
@@ -52,6 +52,8 @@ class Future:
     _loop = None
     _source_traceback = None
     _cancel_message = None
+    # A saved CancelledError for later chaining as an exception context.
+    _cancelled_exc = None
 
     # This field is used for a dual purpose:
     # - Its presence is a marker to declare that a class implements
@@ -124,6 +126,21 @@ class Future:
             raise RuntimeError("Future object is not initialized.")
         return loop
 
+    def _make_cancelled_error(self):
+        """Create the CancelledError to raise if the Future is cancelled.
+
+        This should only be called once when handling a cancellation since
+        it erases the saved context exception value.
+        """
+        if self._cancel_message is None:
+            exc = exceptions.CancelledError()
+        else:
+            exc = exceptions.CancelledError(self._cancel_message)
+        exc.__context__ = self._cancelled_exc
+        # Remove the reference since we don't need this anymore.
+        self._cancelled_exc = None
+        return exc
+
     def cancel(self, msg=None):
         """Cancel the future and schedule callbacks.
 
@@ -175,9 +192,8 @@ class Future:
         the future is done and has an exception set, this exception is raised.
         """
         if self._state == _CANCELLED:
-            raise exceptions.CancelledError(
-                '' if self._cancel_message is None else self._cancel_message)
-
+            exc = self._make_cancelled_error()
+            raise exc
         if self._state != _FINISHED:
             raise exceptions.InvalidStateError('Result is not ready.')
         self.__log_traceback = False
@@ -194,8 +210,8 @@ class Future:
         InvalidStateError.
         """
         if self._state == _CANCELLED:
-            raise exceptions.CancelledError(
-                '' if self._cancel_message is None else self._cancel_message)
+            exc = self._make_cancelled_error()
+            raise exc
         if self._state != _FINISHED:
             raise exceptions.InvalidStateError('Exception is not set.')
         self.__log_traceback = False