Spaces:
Running
Running
| # Adapted with permission from the EdgeDB project; | |
| # license: PSFL. | |
| __all__ = ["TaskGroup"] | |
| from . import events | |
| from . import exceptions | |
| from . import tasks | |
| class TaskGroup: | |
| """Asynchronous context manager for managing groups of tasks. | |
| Example use: | |
| async with asyncio.TaskGroup() as group: | |
| task1 = group.create_task(some_coroutine(...)) | |
| task2 = group.create_task(other_coroutine(...)) | |
| print("Both tasks have completed now.") | |
| All tasks are awaited when the context manager exits. | |
| Any exceptions other than `asyncio.CancelledError` raised within | |
| a task will cancel all remaining tasks and wait for them to exit. | |
| The exceptions are then combined and raised as an `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 | |
| 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 been already entered") | |
| self._entered = True | |
| 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') | |
| return self | |
| async def __aexit__(self, et, exc, tb): | |
| self._exiting = True | |
| if (exc is not None and | |
| self._is_base_error(exc) and | |
| self._base_error is None): | |
| self._base_error = exc | |
| propagate_cancellation_error = \ | |
| exc if et is exceptions.CancelledError else None | |
| if self._parent_cancel_requested: | |
| # If this flag is set we *must* call uncancel(). | |
| if self._parent_task.uncancel() == 0: | |
| # If there are no pending cancellations left, | |
| # don't propagate CancelledError. | |
| propagate_cancellation_error = None | |
| if et is not None: | |
| if not self._aborting: | |
| # Our parent task is being cancelled: | |
| # | |
| # async with TaskGroup() as g: | |
| # g.create_task(...) | |
| # await ... # <- CancelledError | |
| # | |
| # or there's an exception in "async with": | |
| # | |
| # async with TaskGroup() as g: | |
| # g.create_task(...) | |
| # 1 / 0 | |
| # | |
| self._abort() | |
| # We use while-loop here because "self._on_completed_fut" | |
| # can be cancelled multiple times if our parent task | |
| # is being cancelled repeatedly (or even once, when | |
| # our own cancellation is already in progress) | |
| 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: | |
| # Our parent task is being cancelled: | |
| # | |
| # async def wrapper(): | |
| # async with TaskGroup() as g: | |
| # g.create_task(foo) | |
| # | |
| # "wrapper" is being cancelled while "foo" is | |
| # still running. | |
| propagate_cancellation_error = ex | |
| self._abort() | |
| self._on_completed_fut = None | |
| assert not self._tasks | |
| if self._base_error is not None: | |
| raise self._base_error | |
| # Propagate CancelledError if there is one, except if there | |
| # are other errors -- those have priority. | |
| if propagate_cancellation_error and not self._errors: | |
| raise propagate_cancellation_error | |
| if et is not None and et is not exceptions.CancelledError: | |
| self._errors.append(exc) | |
| if self._errors: | |
| # Exceptions are heavy objects that can have object | |
| # cycles (bad for GC); let's not keep a reference to | |
| # a bunch of them. | |
| try: | |
| me = BaseExceptionGroup('unhandled errors in a TaskGroup', self._errors) | |
| raise me from None | |
| finally: | |
| self._errors = None | |
| def create_task(self, coro, *, name=None, context=None): | |
| """Create a new task in this group and return it. | |
| Similar to `asyncio.create_task`. | |
| """ | |
| if not self._entered: | |
| raise RuntimeError(f"TaskGroup {self!r} has not been entered") | |
| if self._exiting and not self._tasks: | |
| raise RuntimeError(f"TaskGroup {self!r} is finished") | |
| if self._aborting: | |
| raise RuntimeError(f"TaskGroup {self!r} is shutting down") | |
| if context is None: | |
| task = self._loop.create_task(coro) | |
| else: | |
| task = self._loop.create_task(coro, context=context) | |
| tasks._set_task_name(task, name) | |
| task.add_done_callback(self._on_task_done) | |
| self._tasks.add(task) | |
| return task | |
| # Since Python 3.8 Tasks propagate all exceptions correctly, | |
| # except for KeyboardInterrupt and SystemExit which are | |
| # still considered special. | |
| 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) | |
| 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(): | |
| # Not sure if this case is possible, but we want to handle | |
| # it anyways. | |
| 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: | |
| # If parent task *is not* being cancelled, it means that we want | |
| # to manually cancel it to abort whatever is being run right now | |
| # in the TaskGroup. But we want to mark parent task as | |
| # "not cancelled" later in __aexit__. Example situation that | |
| # we need to handle: | |
| # | |
| # async def foo(): | |
| # try: | |
| # async with TaskGroup() as g: | |
| # g.create_task(crash_soon()) | |
| # await something # <- this needs to be canceled | |
| # # by the TaskGroup, e.g. | |
| # # foo() needs to be cancelled | |
| # except Exception: | |
| # # Ignore any exceptions raised in the TaskGroup | |
| # pass | |
| # await something_else # this line has to be called | |
| # # after TaskGroup is finished. | |
| self._abort() | |
| self._parent_cancel_requested = True | |
| self._parent_task.cancel() | |