taskgroups.py
python
# Адаптировано с разрешения проекта EdgeDB;
# лицензия: PSFL.
__all__ = ("TaskGroup",)
from . import events
from . import exceptions
from . import futures
from . import tasks
class TaskGroup:
"""Асинхронный контекстный менеджер для управления группами задач.
Пример использования:
async with asyncio.TaskGroup() as group:
task1 = group.create_task(some_coroutine(...))
task2 = group.create_task(other_coroutine(...))
print("Both tasks have completed now.")
Все задачи ожидаются при выходе из контекстного менеджера.
Любые исключения, кроме `asyncio.CancelledError`, возникшие внутри
задачи, отменяют все остальные задачи и ожидают их завершения.
Затем исключения объединяются и вызываются как `ExceptionGroup`.
"""
def __init__(self):
self._entered = False
self._exiting = False
self._aborting = False
self._loop = None
self._parent_task = None
self._parent_cancel_requested = False
self._tasks = set()
self._errors = []
self._base_error = None
self._on_completed_fut = None
self._cancel_on_enter = False
def __repr__(self):
info = ['']
if self._tasks:
info.append(f'tasks={len(self._tasks)}')
if self._errors:
info.append(f'errors={len(self._errors)}')
if self._aborting:
info.append('cancelling')
elif self._entered:
info.append('entered')
info_str = ' '.join(info)
return f'<TaskGroup{info_str}>'
async def __aenter__(self):
if self._entered:
raise RuntimeError(
f"TaskGroup {self!r} has already been entered")
if self._loop is None:
self._loop = events.get_running_loop()
self._parent_task = tasks.current_task(self._loop)
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
async def __aexit__(self, et, exc, tb):
tb = None
try:
return await self._aexit(et, exc)
finally:
# Исключения — это тяжелые объекты, которые могут иметь объектные
# циклы (плохо для сборщика мусора); не будем хранить ссылки на
# их много. Было бы лучше использовать try/finally
# напрямую в __aexit__, но это внесло бы шум в diff
self._parent_task = None
self._errors = None
self._base_error = None
exc = None
async def _aexit(self, et, exc):
self._exiting = True
if (exc is not None and
self._is_base_error(exc) and
self._base_error is None):
self._base_error = exc
if et is not None and issubclass(et, exceptions.CancelledError):
propagate_cancellation_error = exc
else:
propagate_cancellation_error = None
if et is not None:
if not self._aborting:
# Наша родительская задача отменяется:
#
# async with TaskGroup() as g:
# g.create_task(...)
# await ... # <- CancelledError
#
# или есть исключение в "async with":
#
# async with TaskGroup() as g:
# g.create_task(...)
# 1 / 0
#
self._abort()
# Мы используем цикл while здесь, потому что "self._on_completed_fut"
# может быть отменён несколько раз, если наша родительская задача
# отменяется повторно (или даже один раз, когда
# наша собственная отмена уже выполняется)
while self._tasks:
if self._on_completed_fut is None:
self._on_completed_fut = self._loop.create_future()
try:
await self._on_completed_fut
except exceptions.CancelledError as ex:
if not self._aborting:
# Наша родительская задача отменяется:
#
# async def wrapper():
# async with TaskGroup() as g:
# g.create_task(foo)
#
# "wrapper" отменяется, пока "foo"
# всё ещё выполняется.
propagate_cancellation_error = ex
self._abort()
self._on_completed_fut = None
assert not self._tasks
if self._base_error is not None:
try:
raise self._base_error
finally:
exc = None
if self._parent_cancel_requested:
# Если этот флаг установлен, мы *должны* вызвать uncancel().
if self._parent_task.uncancel() == 0:
# Если не осталось ожидающих отмен,
# не распространяйте CancelledError.
propagate_cancellation_error = None
# Распространить CancelledError, если он есть, за исключением случаев, когда
# имеются другие ошибки — они имеют приоритет.
try:
if propagate_cancellation_error is not None and not self._errors:
try:
raise propagate_cancellation_error
finally:
exc = None
finally:
propagate_cancellation_error = None
if et is not None and not issubclass(et, exceptions.CancelledError):
self._errors.append(exc)
if self._errors:
# Если родительская задача отменяется извне
# группы задач, отменить отмену и повторно отменить родительскую задачу,
# что сохранит стабильным счетчик отмен.
if self._parent_task.cancelling():
self._parent_task.uncancel()
self._parent_task.cancel()
try:
raise BaseExceptionGroup(
'unhandled errors in a TaskGroup',
self._errors,
) from None
finally:
exc = None
# Подавить любые оставшиеся исключения (исключения, которые заслуживают быть вызванными
# были вызваны выше).
return True
def create_task(self, coro, **kwargs):
"""Создать новую задачу в этой группе и вернуть её.
Аналог `asyncio.create_task`.
"""
if not self._entered:
coro.close()
raise RuntimeError(f"TaskGroup {self!r} has not been entered")
if self._exiting and not self._tasks:
coro.close()
raise RuntimeError(f"TaskGroup {self!r} is finished")
if self._aborting:
coro.close()
raise RuntimeError(f"TaskGroup {self!r} is shutting down")
task = self._loop.create_task(coro, **kwargs)
futures.future_add_to_awaited_by(task, self._parent_task)
# Всегда планируйте колбэк done, даже если задача
# уже выполнена (например, если корутине удалось завершиться нетерпеливо),
# иначе если задача завершится с исключением, она отменит
# текущую задачу слишком рано. gh-128550, gh-128588
self._tasks.add(task)
task.add_done_callback(self._on_task_done)
try:
return task
finally:
# gh-128552: предотвратить циклическую ссылку
# task.exception().__traceback__->TaskGroup.create_task->task
del task
# Начиная с Python 3.8 задачи корректно распространяют все исключения,
# за исключением KeyboardInterrupt и SystemExit, которые
# по-прежнему считаются особенными.
def _is_base_error(self, exc: BaseException) -> bool:
assert isinstance(exc, BaseException)
return isinstance(exc, (SystemExit, KeyboardInterrupt))
def _abort(self):
self._aborting = True
for t in self._tasks:
if not t.done():
t.cancel()
def _on_task_done(self, task):
self._tasks.discard(task)
futures.future_discard_from_awaited_by(task, self._parent_task)
if self._on_completed_fut is not None and not self._tasks:
if not self._on_completed_fut.done():
self._on_completed_fut.set_result(True)
if task.cancelled():
return
exc = task.exception()
if exc is None:
return
self._errors.append(exc)
if self._is_base_error(exc) and self._base_error is None:
self._base_error = exc
if self._parent_task.done():
# Не уверены, что такой случай возможен, но мы хотим обработать
# его в любом случае.
self._loop.call_exception_handler({
'message': f'Task {task!r} has errored out but its parent '
f'task {self._parent_task} is already completed',
'exception': exc,
'task': task,
})
return
if not self._aborting and not self._parent_cancel_requested:
# Если родительская задача *не* отменяется, это означает, что мы хотим
# отменить её вручную, чтобы прервать то, что сейчас выполняется
# в TaskGroup. Но мы хотим пометить родительскую задачу как
# "не отменённая" позже в __aexit__. Пример ситуации, которую
# нам нужно обработать:
#
# async def foo():
# try:
# async with TaskGroup() as g:
# g.create_task(crash_soon())
# await something # <- эту нужно отменить
# # с помощью TaskGroup, например,
# # foo() нужно отменить
# except Exception:
# # Игнорировать любые исключения, вызванные в TaskGroup
# pass
# await something_else # эта строка должна быть вызвана
# # после завершения TaskGroup.
self._abort()
self._parent_cancel_requested = True
self._parent_task.cancel()
def cancel(self):
"""Отменить группу задач
`cancel()` будет вызван для всех задач в группе, которые ещё не завершены,
а также для родительской (основной) задачи группы. Это приведёт к тому,
что менеджер контекста группы задач завершится *без* возбуждения
`asyncio.CancelledError`.
Если `cancel()` вызывается до входа в группу задач, группа будет
отменена при входе. Это полезно для сценариев, когда один фрагмент кода
передаёт неиспользуемый экземпляр TaskGroup другому, чтобы иметь
возможность отменить всё, что выполняется в группе.
`cancel()` идемпотентен и может быть вызван после того, как группа задач
уже завершилась.
"""
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()