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."
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 @@
_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 @@
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 @@
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 @@
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
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index a3a0a33..21b98b6 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -270,8 +270,7 @@
f'_step(): already done: {self!r}, {exc!r}')
if self._must_cancel:
if not isinstance(exc, exceptions.CancelledError):
- exc = exceptions.CancelledError(''
- if self._cancel_message is None else self._cancel_message)
+ exc = self._make_cancelled_error()
self._must_cancel = False
coro = self._coro
self._fut_waiter = None
@@ -293,11 +292,9 @@
else:
super().set_result(exc.value)
except exceptions.CancelledError as exc:
- if exc.args:
- cancel_msg = exc.args[0]
- else:
- cancel_msg = None
- super().cancel(msg=cancel_msg) # I.e., Future.cancel(self).
+ # Save the original exception so we can chain it later.
+ self._cancelled_exc = exc
+ super().cancel() # I.e., Future.cancel(self).
except (KeyboardInterrupt, SystemExit) as exc:
super().set_exception(exc)
raise
@@ -787,8 +784,7 @@
# Check if 'fut' is cancelled first, as
# 'fut.exception()' will *raise* a CancelledError
# instead of returning it.
- exc = exceptions.CancelledError(''
- if fut._cancel_message is None else fut._cancel_message)
+ exc = fut._make_cancelled_error()
outer.set_exception(exc)
return
else:
@@ -804,9 +800,12 @@
for fut in children:
if fut.cancelled():
- # Check if 'fut' is cancelled first, as
- # 'fut.exception()' will *raise* a CancelledError
- # instead of returning it.
+ # Check if 'fut' is cancelled first, as 'fut.exception()'
+ # will *raise* a CancelledError instead of returning it.
+ # Also, since we're adding the exception return value
+ # to 'results' instead of raising it, don't bother
+ # setting __context__. This also lets us preserve
+ # calling '_make_cancelled_error()' at most once.
res = exceptions.CancelledError(
'' if fut._cancel_message is None else
fut._cancel_message)
@@ -820,8 +819,7 @@
# If gather is being cancelled we must propagate the
# cancellation regardless of *return_exceptions* argument.
# See issue 32684.
- exc = exceptions.CancelledError(''
- if fut._cancel_message is None else fut._cancel_message)
+ exc = fut._make_cancelled_error()
outer.set_exception(exc)
else:
outer.set_result(results)