Coverage for Lib/asyncio/taskgroups.py: 97%
161 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-25 03:29 +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 pending_cancellation_error = None
120 while self._tasks:
121 if self._on_completed_fut is None: 121 ↛ 124line 121 didn't jump to line 124 because the condition on line 121 was always true
122 self._on_completed_fut = self._loop.create_future()
124 try:
125 await self._on_completed_fut
126 except exceptions.CancelledError as ex:
127 pending_cancellation_error = ex
128 if not self._aborting:
129 # Our parent task is being cancelled:
130 #
131 # async def wrapper():
132 # async with TaskGroup() as g:
133 # g.create_task(foo)
134 #
135 # "wrapper" is being cancelled while "foo" is
136 # still running.
137 propagate_cancellation_error = ex
138 self._abort()
140 self._on_completed_fut = None
142 assert not self._tasks
144 if self._base_error is not None:
145 # self._base_error (SystemExit or KeyboardInterrupt) is about
146 # to propagate out of this method, which discards any other
147 # collected task errors silently. Report them instead of
148 # losing them. See gh-135736.
149 for suppressed_exc in self._errors:
150 self._loop.call_exception_handler({
151 'message': 'TaskGroup task exception was not '
152 'propagated because the TaskGroup body '
153 'is being closed with a BaseException',
154 'exception': suppressed_exc,
155 'task_group': self,
156 })
157 try:
158 raise self._base_error
159 finally:
160 exc = None
162 if self._parent_cancel_requested:
163 # If this flag is set we *must* call uncancel().
164 if self._parent_task.uncancel() == 0:
165 # If there are no pending cancellations left,
166 # don't propagate CancelledError.
167 propagate_cancellation_error = None
168 elif propagate_cancellation_error is None:
169 # gh-155433: the remaining cancellation is not ours, don't drop it
170 propagate_cancellation_error = pending_cancellation_error
172 # Propagate CancelledError if there is one, except if there
173 # are other errors -- those have priority.
174 try:
175 if propagate_cancellation_error is not None and not self._errors:
176 try:
177 raise propagate_cancellation_error
178 finally:
179 exc = None
180 finally:
181 propagate_cancellation_error = None
183 if et is not None and not issubclass(et, exceptions.CancelledError):
184 self._errors.append(exc)
186 if self._errors:
187 # If the parent task is being cancelled from the outside
188 # of the taskgroup, un-cancel and re-cancel the parent task,
189 # which will keep the cancel count stable.
190 if self._parent_task.cancelling():
191 self._parent_task.uncancel()
192 self._parent_task.cancel()
193 try:
194 # If the *only* error is a GeneratorExit from the body
195 # of the group, then instead of raising an
196 # ExceptionGroup we raise GeneratorExit. This ensures
197 # that async generators that use TaskGroup properly
198 # swallow the exception on `aclose()` while ensuring
199 # that no exceptions from subtasks are swallowed.
200 if (
201 et is not None
202 and issubclass(et, GeneratorExit)
203 and len(self._errors) == 1
204 ):
205 raise exc
206 else:
207 raise BaseExceptionGroup(
208 'unhandled errors in a TaskGroup',
209 self._errors,
210 ) from None
211 finally:
212 exc = None
214 # Suppress any remaining exception (exceptions deserving to be raised
215 # were raised above).
216 return True
218 def create_task(self, coro, **kwargs):
219 """Create a new task in this group and return it.
221 Similar to `asyncio.create_task`.
222 """
223 if not self._entered:
224 coro.close()
225 raise RuntimeError(f"TaskGroup {self!r} has not been entered")
226 if self._exiting and not self._tasks:
227 coro.close()
228 raise RuntimeError(f"TaskGroup {self!r} is finished")
229 if self._aborting:
230 coro.close()
231 raise RuntimeError(f"TaskGroup {self!r} is shutting down")
232 task = self._loop.create_task(coro, **kwargs)
234 futures.future_add_to_awaited_by(task, self._parent_task)
236 # Always schedule the done callback even if the task is
237 # already done (e.g. if the coro was able to complete eagerly),
238 # otherwise if the task completes with an exception then it will cancel
239 # the current task too early. gh-128550, gh-128588
240 self._tasks.add(task)
241 task.add_done_callback(self._on_task_done)
242 # gh-155418: an eager task can cancel the group before joining _tasks
243 if self._aborting and not task.done():
244 task.cancel()
245 try:
246 return task
247 finally:
248 # gh-128552: prevent a refcycle of
249 # task.exception().__traceback__->TaskGroup.create_task->task
250 del task
252 # Since Python 3.8 Tasks propagate all exceptions correctly,
253 # except for KeyboardInterrupt and SystemExit which are
254 # still considered special.
256 def _is_base_error(self, exc: BaseException) -> bool:
257 assert isinstance(exc, BaseException)
258 return isinstance(exc, (SystemExit, KeyboardInterrupt))
260 def _abort(self):
261 self._aborting = True
263 for t in self._tasks:
264 if not t.done():
265 t.cancel()
267 def _on_task_done(self, task):
268 self._tasks.discard(task)
270 futures.future_discard_from_awaited_by(task, self._parent_task)
272 if self._on_completed_fut is not None and not self._tasks:
273 if not self._on_completed_fut.done():
274 self._on_completed_fut.set_result(True)
276 if task.cancelled():
277 return
279 exc = task.exception()
280 if exc is None:
281 return
283 self._errors.append(exc)
284 if self._is_base_error(exc) and self._base_error is None: 284 ↛ 285line 284 didn't jump to line 285 because the condition on line 284 was never true
285 self._base_error = exc
287 if self._parent_task.done(): 287 ↛ 290line 287 didn't jump to line 290 because the condition on line 287 was never true
288 # Not sure if this case is possible, but we want to handle
289 # it anyways.
290 self._loop.call_exception_handler({
291 'message': f'Task {task!r} has errored out but its parent '
292 f'task {self._parent_task} is already completed',
293 'exception': exc,
294 'task': task,
295 })
296 return
298 if not self._aborting and not self._parent_cancel_requested:
299 # If parent task *is not* being cancelled, it means that we want
300 # to manually cancel it to abort whatever is being run right now
301 # in the TaskGroup. But we want to mark parent task as
302 # "not cancelled" later in __aexit__. Example situation that
303 # we need to handle:
304 #
305 # async def foo():
306 # try:
307 # async with TaskGroup() as g:
308 # g.create_task(crash_soon())
309 # await something # <- this needs to be canceled
310 # # by the TaskGroup, e.g.
311 # # foo() needs to be cancelled
312 # except Exception:
313 # # Ignore any exceptions raised in the TaskGroup
314 # pass
315 # await something_else # this line has to be called
316 # # after TaskGroup is finished.
317 self._abort()
318 self._parent_cancel_requested = True
319 self._parent_task.cancel()
321 def cancel(self):
322 """Cancel the task group
324 `cancel()` will be called on any tasks in the group that aren't yet
325 done, as well as the parent (body) of the group. This will cause
326 the task group context manager to exit *without*
327 `asyncio.CancelledError` being raised.
329 If `cancel()` is called before entering the task group, the group
330 will be cancelled upon entry. This is useful for patterns where
331 one piece of code passes an unused TaskGroup instance to another in
332 order to have the ability to cancel anything run within the group.
334 `cancel()` is idempotent and may be called after the task group has
335 already exited.
336 """
337 if not self._entered:
338 self._cancel_on_enter = True
339 return
340 if self._exiting and not self._tasks:
341 return
342 if not self._aborting:
343 self._abort()
344 if self._parent_task and not self._parent_cancel_requested: 344 ↛ exitline 344 didn't return from function 'cancel' because the condition on line 344 was always true
345 self._parent_cancel_requested = True
346 self._parent_task.cancel()