diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 571013745aa03aa..3c4b2e265869a45 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -4,7 +4,7 @@ 'Task', 'create_task', 'FIRST_COMPLETED', 'FIRST_EXCEPTION', 'ALL_COMPLETED', 'wait', 'wait_for', 'as_completed', 'sleep', - 'gather', 'shield', 'ensure_future', 'run_coroutine_threadsafe', + 'gather', 'shield', 'shield_scope', 'ensure_future', 'run_coroutine_threadsafe', 'current_task', 'all_tasks', '_register_task', '_unregister_task', '_enter_task', '_leave_task', ) @@ -17,6 +17,7 @@ import types import warnings import weakref +import contextlib from types import GenericAlias from . import base_tasks @@ -75,6 +76,20 @@ def _set_task_name(task, name): set_name(name) +@contextlib.contextmanager +def shield_scope(): + task = current_task() + task.shields += 1 + try: + yield + finally: + task.shields -= 1 + if task.shields < 0: + raise RuntimeError( + f"Unbalanced enter/exit of shield scopes for task {task!r}" + ) + + class Task(futures._PyFuture): # Inherit Python Task implementation # from a Python Future implementation. @@ -95,6 +110,7 @@ class Task(futures._PyFuture): # Inherit Python Task implementation def __init__(self, coro, *, loop=None, name=None, context=None): super().__init__(loop=loop) + self.shields = 0 if self._source_traceback: del self._source_traceback[-1] if not coroutines.iscoroutine(coro): @@ -219,7 +235,7 @@ def cancel(self, msg=None): # Also remember that this is duplicated in _asynciomodule.c. # if self._num_cancels_requested > 1: # return False - if self._fut_waiter is not None: + if self._fut_waiter is not None and not self.shields: if self._fut_waiter.cancel(msg=msg): # Leave self._fut_waiter; it may be a Task that # catches and ignores the cancellation so we may have @@ -254,7 +270,7 @@ def __step(self, exc=None): if self.done(): raise exceptions.InvalidStateError( f'_step(): already done: {self!r}, {exc!r}') - if self._must_cancel: + if self._must_cancel and not self.shields: if not isinstance(exc, exceptions.CancelledError): exc = self._make_cancelled_error() self._must_cancel = False @@ -307,7 +323,7 @@ def __step(self, exc=None): result.add_done_callback( self.__wakeup, context=self._context) self._fut_waiter = result - if self._must_cancel: + if self._must_cancel and not self.shields: if self._fut_waiter.cancel( msg=self._cancel_message): self._must_cancel = False diff --git a/Misc/NEWS.d/next/Library/2022-11-30-09-00-13.gh-issue-99714.Xkbibc.rst b/Misc/NEWS.d/next/Library/2022-11-30-09-00-13.gh-issue-99714.Xkbibc.rst new file mode 100644 index 000000000000000..76f6cdea13c71f8 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2022-11-30-09-00-13.gh-issue-99714.Xkbibc.rst @@ -0,0 +1,3 @@ +Add :func:`shield_scope` context manager in :mod:`asyncio.tasks` to shield +async blocks of code from cancellation. Add :func:`shielded` decorator in +:mod:`asyncio.tasks` to shield functions from cancellation.