add TaskGroup.stop()

ISSUE: #108951

Author: belm0

Doc/library/asyncio-task.rst

@@ -342,6 +342,25 @@ and reliable way to wait for all tasks in the group to finish.
 
          Close the given coroutine if the task group is not active.
 
+   .. method:: stop()
+
+      Stop the task group
+
+      :meth:`~asyncio.Task.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* a
+      :exc:`asyncio.CancelledError` being raised.
+
+      If :meth:`stop` is called before entering the task group, the group will be
+      stopped 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 stop anything run within the group.
+
+      :meth:`stop` is idempotent and may be called after the task group has
+      already exited.
+
+      .. versionadded:: 3.14
+
 Example::
 
     async def main():
@@ -414,53 +433,6 @@ reported by :meth:`asyncio.Task.cancelling`.
    Improved handling of simultaneous internal and external cancellations
    and correct preservation of cancellation counts.
 
-Terminating a Task Group
-------------------------
-
-While terminating a task group is not natively supported by the standard
-library, termination can be achieved by adding an exception-raising task
-to the task group and ignoring the raised exception:
-
-.. code-block:: python
-
-   import asyncio
-   from asyncio import TaskGroup
-
-   class TerminateTaskGroup(Exception):
-       """Exception raised to terminate a task group."""
-
-   async def force_terminate_task_group():
-       """Used to force termination of a task group."""
-       raise TerminateTaskGroup()
-
-   async def job(task_id, sleep_time):
-       print(f'Task {task_id}: start')
-       await asyncio.sleep(sleep_time)
-       print(f'Task {task_id}: done')
-
-   async def main():
-       try:
-           async with TaskGroup() as group:
-               # spawn some tasks
-               group.create_task(job(1, 0.5))
-               group.create_task(job(2, 1.5))
-               # sleep for 1 second
-               await asyncio.sleep(1)
-               # add an exception-raising task to force the group to terminate
-               group.create_task(force_terminate_task_group())
-       except* TerminateTaskGroup:
-           pass
-
-   asyncio.run(main())
-
-Expected output:
-
-.. code-block:: text
-
-   Task 1: start
-   Task 2: start
-   Task 1: done
-
 Sleeping
 ========
 

Lib/asyncio/taskgroups.py

@@ -36,6 +36,7 @@ def __init__(self):
         self._errors = []
         self._base_error = None
         self._on_completed_fut = None
+        self._stop_on_enter = False
 
     def __repr__(self):
         info = ['']
@@ -62,6 +63,8 @@ async def __aenter__(self):
             raise RuntimeError(
                 f'TaskGroup {self!r} cannot determine the parent task')
         self._entered = True
+        if self._stop_on_enter:
+            self.stop()
 
         return self
 
@@ -147,6 +150,10 @@ async def _aexit(self, et, exc):
                 # If there are no pending cancellations left,
                 # don't propagate CancelledError.
                 propagate_cancellation_error = None
+                # If Cancelled would actually be raised out of the TaskGroup,
+                # suppress it-- this is significant when using stop().
+                if not self._errors:
+                    return True
 
         # Propagate CancelledError if there is one, except if there
         # are other errors -- those have priority.
@@ -273,3 +280,30 @@ def _on_task_done(self, task):
             self._abort()
             self._parent_cancel_requested = True
             self._parent_task.cancel()
+
+    def stop(self):
+        """Stop 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* a Cancelled exception
+        being raised.
+
+        If `stop()` is called before entering the task group, the group will be
+        stopped 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 stop anything run within the group.
+
+        `stop()` is idempotent and may be called after the task group has
+        already exited.
+        """
+        if not self._entered:
+            self._stop_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()

Lib/test/test_asyncio/test_taskgroups.py

@@ -3,10 +3,12 @@
 
 import sys
 import gc
+
 import asyncio
 import contextvars
 import contextlib
 from asyncio import taskgroups
+import math
 import unittest
 import warnings
 
@@ -997,6 +999,69 @@ class MyKeyboardInterrupt(KeyboardInterrupt):
         self.assertIsNotNone(exc)
         self.assertListEqual(gc.get_referrers(exc), no_other_refs())
 
+    async def test_taskgroup_stop_children(self):
+        async with asyncio.TaskGroup() as tg:
+            tg.create_task(asyncio.sleep(math.inf))
+            tg.create_task(asyncio.sleep(math.inf))
+            await asyncio.sleep(0)
+            tg.stop()
+
+    async def test_taskgroup_stop_body(self):
+        count = 0
+        async with asyncio.TaskGroup() as tg:
+            tg.stop()
+            count += 1
+            await asyncio.sleep(0)
+            count += 1
+        self.assertEqual(count, 1)
+
+    async def test_taskgroup_stop_idempotent(self):
+        count = 0
+        async with asyncio.TaskGroup() as tg:
+            tg.stop()
+            tg.stop()
+            count += 1
+            await asyncio.sleep(0)
+            count += 1
+        self.assertEqual(count, 1)
+
+    async def test_taskgroup_stop_after_exit(self):
+        async with asyncio.TaskGroup() as tg:
+            await asyncio.sleep(0)
+        tg.stop()
+
+    async def test_taskgroup_stop_before_enter(self):
+        tg = asyncio.TaskGroup()
+        tg.stop()
+        count = 0
+        async with tg:
+            count += 1
+            await asyncio.sleep(0)
+            count += 1
+        self.assertEqual(count, 1)
+
+    async def test_taskgroup_stop_before_exception(self):
+        async def raise_exc(parent_tg: asyncio.TaskGroup):
+            parent_tg.stop()
+            raise RuntimeError
+
+        with self.assertRaises(ExceptionGroup):
+            async with asyncio.TaskGroup() as tg:
+                tg.create_task(raise_exc(tg))
+                await asyncio.sleep(1)
+
+    async def test_taskgroup_stop_after_exception(self):
+        async def raise_exc(parent_tg: asyncio.TaskGroup):
+            try:
+                raise RuntimeError
+            finally:
+                parent_tg.stop()
+
+        with self.assertRaises(ExceptionGroup):
+            async with asyncio.TaskGroup() as tg:
+                tg.create_task(raise_exc(tg))
+                await asyncio.sleep(1)
+
 
 if __name__ == "__main__":
     unittest.main()