diff --git a/taskgroup/taskgroups.py b/taskgroup/taskgroups.py index 212c5fd..f884984 100644 --- a/taskgroup/taskgroups.py +++ b/taskgroup/taskgroups.py @@ -82,6 +82,7 @@ def __init__(self) -> None: self._errors: list[BaseException] | None = [] self._base_error: BaseException | None = None self._on_completed_fut: futures.Future[Literal[True]] | None = None + self._cancel_on_enter = False def __repr__(self) -> str: info = [""] @@ -106,6 +107,8 @@ async def __aenter__(self) -> Self: if self._parent_task is None: raise RuntimeError(f"TaskGroup {self!r} cannot determine the parent task") self._entered = True + if self._cancel_on_enter: + self.cancel() return self @@ -222,6 +225,10 @@ async def _aexit( finally: exc = None + # Suppress any remaining exception (exceptions deserving to be raised + # were raised above). + return True + def create_task( self, coro: _TaskCompatibleCoro[_T_co], @@ -327,6 +334,33 @@ def _on_task_done(self, task: tasks.Task[object]) -> None: self._parent_cancel_requested = True self._parent_task.cancel() + def cancel(self): + """Cancel the task group + + `cancel()` will be called on any tasks in the group that aren't yet + done, as well as the parent (body) of the group. This will cause the + task group context manager to exit *without* `asyncio.CancelledError` + being raised. + + If `cancel()` is called before entering the task group, the group will be + cancelled upon entry. This is useful for patterns where one piece of + code passes an unused TaskGroup instance to another in order to have + the ability to cancel anything run within the group. + + `cancel()` is idempotent and may be called after the task group has + already exited. + """ + if not self._entered: + self._cancel_on_enter = True + return + if self._exiting and not self._tasks: + return + if not self._aborting: + self._abort() + if self._parent_task and not self._parent_cancel_requested: + self._parent_cancel_requested = True + self._parent_task.cancel() + class TaskGroup(_TaskGroup): __stack: contextlib.AsyncExitStack