Coverage for Lib/asyncio/taskgroups.py: 97%
153 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 01:31 +0000
1# Adapted with permission from the EdgeDB project;
2# license: PSFL.
5__all__ = ("TaskGroup",)
7from . import events
8from . import exceptions
9from . import futures
10from . import tasks
13class TaskGroup:
14 """Asynchronous context manager for managing groups of tasks.
16 Example use:
18 async with asyncio.TaskGroup() as group:
19 task1 = group.create_task(some_coroutine(...))
20 task2 = group.create_task(other_coroutine(...))
21 print("Both tasks have completed now.")
23 All tasks are awaited when the context manager exits.
25 Any exceptions other than `asyncio.CancelledError` raised within
26 a task will cancel all remaining tasks and wait for them to exit.
27 The exceptions are then combined and raised as an `ExceptionGroup`.
28 """
29 def __init__(self):
30 self._entered = False
31 self._exiting = False
32 self._aborting = False
33 self._loop = None
34 self._parent_task = None
35 self._parent_cancel_requested = False
36 self._tasks = set()
37 self._errors = []
38 self._base_error = None
39 self._on_completed_fut = None
40 self._cancel_on_enter = False
42 def __repr__(self):
43 info = ['']
44 if self._tasks:
45 info.append(f'tasks={len(self._tasks)}')
46 if self._errors:
47 info.append(f'errors={len(self._errors)}')
48 if self._aborting:
49 info.append('cancelling')
50 elif self._entered:
51 info.append('entered')
53 info_str = ' '.join(info)
54 return f'<TaskGroup{info_str}>'
56 async def __aenter__(self):
57 if self._entered:
58 raise RuntimeError(
59 f"TaskGroup {self!r} has already been entered")
60 if self._loop is None: 60 ↛ 62line 60 didn't jump to line 62 because the condition on line 60 was always true
61 self._loop = events.get_running_loop()
62 self._parent_task = tasks.current_task(self._loop)
63 if self._parent_task is None:
64 raise RuntimeError(
65 f'TaskGroup {self!r} cannot determine the parent task')
66 self._entered = True
67 if self._cancel_on_enter:
68 self.cancel()
70 return self
72 async def __aexit__(self, et, exc, tb):
73 tb = None
74 try:
75 return await self._aexit(et, exc)
76 finally:
77 # Exceptions are heavy objects that can have object
78 # cycles (bad for GC); let's not keep a reference to
79 # a bunch of them. It would be nicer to use a try/finally
80 # in __aexit__ directly but that introduced some diff noise
81 self._parent_task = None
82 self._errors = None
83 self._base_error = None
84 exc = None
86 async def _aexit(self, et, exc):
87 self._exiting = True
89 if (exc is not None and
90 self._is_base_error(exc) and
91 self._base_error is None):
92 self._base_error = exc
94 if et is not None and issubclass(et, exceptions.CancelledError):
95 propagate_cancellation_error = exc
96 else:
97 propagate_cancellation_error = None
99 if et is not None:
100 if not self._aborting:
101 # Our parent task is being cancelled:
102 #
103 # async with TaskGroup() as g:
104 # g.create_task(...)
105 # await ... # <- CancelledError
106 #
107 # or there's an exception in "async with":
108 #
109 # async with TaskGroup() as g:
110 # g.create_task(...)
111 # 1 / 0
112 #
113 self._abort()
115 # We use while-loop here because "self._on_completed_fut"
116 # can be cancelled multiple times if our parent task
117 # is being cancelled repeatedly (or even once, when
118 # our own cancellation is already in progress)
119 while self._tasks:
120 if self._on_completed_fut is None: 120 ↛ 123line 120 didn't jump to line 123 because the condition on line 120 was always true
121 self._on_completed_fut = self._loop.create_future()
123 try:
124 await self._on_completed_fut
125 except exceptions.CancelledError as ex:
126 if not self._aborting:
127 # Our parent task is being cancelled:
128 #
129 # async def wrapper():
130 # async with TaskGroup() as g:
131 # g.create_task(foo)
132 #
133 # "wrapper" is being cancelled while "foo" is
134 # still running.
135 propagate_cancellation_error = ex
136 self._abort()
138 self._on_completed_fut = None
140 assert not self._tasks
142 if self._base_error is not None:
143 try:
144 raise self._base_error
145 finally:
146 exc = None
148 if self._parent_cancel_requested:
149 # If this flag is set we *must* call uncancel().
150 if self._parent_task.uncancel() == 0:
151 # If there are no pending cancellations left,
152 # don't propagate CancelledError.
153 propagate_cancellation_error = None
155 # Propagate CancelledError if there is one, except if there
156 # are other errors -- those have priority.
157 try:
158 if propagate_cancellation_error is not None and not self._errors:
159 try:
160 raise propagate_cancellation_error
161 finally:
162 exc = None
163 finally:
164 propagate_cancellation_error = None
166 if et is not None and not issubclass(et, exceptions.CancelledError):
167 self._errors.append(exc)
169 if self._errors:
170 # If the parent task is being cancelled from the outside
171 # of the taskgroup, un-cancel and re-cancel the parent task,
172 # which will keep the cancel count stable.
173 if self._parent_task.cancelling():
174 self._parent_task.uncancel()
175 self._parent_task.cancel()
176 try:
177 # If the *only* error is a GeneratorExit from the body
178 # of the group, then instead of raising an
179 # ExceptionGroup we raise GeneratorExit. This ensures
180 # that async generators that use TaskGroup properly
181 # swallow the exception on `aclose()` while ensuring
182 # that no exceptions from subtasks are swallowed.
183 if (
184 et is not None
185 and issubclass(et, GeneratorExit)
186 and len(self._errors) == 1
187 ):
188 raise exc
189 else:
190 raise BaseExceptionGroup(
191 'unhandled errors in a TaskGroup',
192 self._errors,
193 ) from None
194 finally:
195 exc = None
197 # Suppress any remaining exception (exceptions deserving to be raised
198 # were raised above).
199 return True
201 def create_task(self, coro, **kwargs):
202 """Create a new task in this group and return it.
204 Similar to `asyncio.create_task`.
205 """
206 if not self._entered:
207 coro.close()
208 raise RuntimeError(f"TaskGroup {self!r} has not been entered")
209 if self._exiting and not self._tasks:
210 coro.close()
211 raise RuntimeError(f"TaskGroup {self!r} is finished")
212 if self._aborting:
213 coro.close()
214 raise RuntimeError(f"TaskGroup {self!r} is shutting down")
215 task = self._loop.create_task(coro, **kwargs)
217 futures.future_add_to_awaited_by(task, self._parent_task)
219 # Always schedule the done callback even if the task is
220 # already done (e.g. if the coro was able to complete eagerly),
221 # otherwise if the task completes with an exception then it will cancel
222 # the current task too early. gh-128550, gh-128588
223 self._tasks.add(task)
224 task.add_done_callback(self._on_task_done)
225 try:
226 return task
227 finally:
228 # gh-128552: prevent a refcycle of
229 # task.exception().__traceback__->TaskGroup.create_task->task
230 del task
232 # Since Python 3.8 Tasks propagate all exceptions correctly,
233 # except for KeyboardInterrupt and SystemExit which are
234 # still considered special.
236 def _is_base_error(self, exc: BaseException) -> bool:
237 assert isinstance(exc, BaseException)
238 return isinstance(exc, (SystemExit, KeyboardInterrupt))
240 def _abort(self):
241 self._aborting = True
243 for t in self._tasks:
244 if not t.done():
245 t.cancel()
247 def _on_task_done(self, task):
248 self._tasks.discard(task)
250 futures.future_discard_from_awaited_by(task, self._parent_task)
252 if self._on_completed_fut is not None and not self._tasks:
253 if not self._on_completed_fut.done():
254 self._on_completed_fut.set_result(True)
256 if task.cancelled():
257 return
259 exc = task.exception()
260 if exc is None:
261 return
263 self._errors.append(exc)
264 if self._is_base_error(exc) and self._base_error is None: 264 ↛ 265line 264 didn't jump to line 265 because the condition on line 264 was never true
265 self._base_error = exc
267 if self._parent_task.done(): 267 ↛ 270line 267 didn't jump to line 270 because the condition on line 267 was never true
268 # Not sure if this case is possible, but we want to handle
269 # it anyways.
270 self._loop.call_exception_handler({
271 'message': f'Task {task!r} has errored out but its parent '
272 f'task {self._parent_task} is already completed',
273 'exception': exc,
274 'task': task,
275 })
276 return
278 if not self._aborting and not self._parent_cancel_requested:
279 # If parent task *is not* being cancelled, it means that we want
280 # to manually cancel it to abort whatever is being run right now
281 # in the TaskGroup. But we want to mark parent task as
282 # "not cancelled" later in __aexit__. Example situation that
283 # we need to handle:
284 #
285 # async def foo():
286 # try:
287 # async with TaskGroup() as g:
288 # g.create_task(crash_soon())
289 # await something # <- this needs to be canceled
290 # # by the TaskGroup, e.g.
291 # # foo() needs to be cancelled
292 # except Exception:
293 # # Ignore any exceptions raised in the TaskGroup
294 # pass
295 # await something_else # this line has to be called
296 # # after TaskGroup is finished.
297 self._abort()
298 self._parent_cancel_requested = True
299 self._parent_task.cancel()
301 def cancel(self):
302 """Cancel the task group
304 `cancel()` will be called on any tasks in the group that aren't yet
305 done, as well as the parent (body) of the group. This will cause
306 the task group context manager to exit *without*
307 `asyncio.CancelledError` being raised.
309 If `cancel()` is called before entering the task group, the group
310 will be cancelled upon entry. This is useful for patterns where
311 one piece of code passes an unused TaskGroup instance to another in
312 order to have the ability to cancel anything run within the group.
314 `cancel()` is idempotent and may be called after the task group has
315 already exited.
316 """
317 if not self._entered:
318 self._cancel_on_enter = True
319 return
320 if self._exiting and not self._tasks:
321 return
322 if not self._aborting:
323 self._abort()
324 if self._parent_task and not self._parent_cancel_requested: 324 ↛ exitline 324 didn't return from function 'cancel' because the condition on line 324 was always true
325 self._parent_cancel_requested = True
326 self._parent_task.cancel()