adapted with permission from the edgedb project license psfl asynchronous context manager for managing groups of tasks example use async with asyncio taskgroup as group task1 group createtasksomecoroutine task2 group createtaskothercoroutine printboth 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 if this flag is set we must call uncancel if there are no pending cancellations left don t propagate cancellederror our parent task is being cancelled async with taskgroup as g g createtask await cancellederror or there s an exception in async with async with taskgroup as g g createtask 1 0 we use whileloop here because self oncompletedfut can be cancelled multiple times if our parent task is being cancelled repeatedly or even once when our own cancellation is already in progress our parent task is being cancelled async def wrapper async with taskgroup as g g createtaskfoo wrapper is being cancelled while foo is still running propagate cancellederror if there is one except if there are other errors those have priority exceptions are heavy objects that can have object cycles bad for gc let s not keep a reference to a bunch of them create a new task in this group and return it similar to asyncio createtask optimization immediately call the done callback if the task is already done e g if the coro was able to complete eagerly and skip scheduling a done callback since python 3 8 tasks propagate all exceptions correctly except for keyboardinterrupt and systemexit which are still considered special not sure if this case is possible but we want to handle it anyways 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 createtaskcrashsoon 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 somethingelse this line has to be called after taskgroup is finished adapted with permission from the edgedb project license psfl 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 if this flag is set we must call uncancel if there are no pending cancellations left don t propagate cancellederror 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 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 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 cancellederror if there is one except if there are other errors those have priority exceptions are heavy objects that can have object cycles bad for gc let s not keep a reference to a bunch of them create a new task in this group and return it similar to asyncio create_task optimization immediately call the done callback if the task is already done e g if the coro was able to complete eagerly and skip scheduling a done callback since python 3 8 tasks propagate all exceptions correctly except for keyboardinterrupt and systemexit which are still considered special not sure if this case is possible but we want to handle it anyways 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
__all__ = ("TaskGroup",)

from . import events
from . import exceptions
from . import tasks


class TaskGroup:
    
    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 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

        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 self._parent_task.uncancel() == 0:
                
                
                propagate_cancellation_error = None

        if et is not None:
            if not self._aborting:
                
                self._abort()

    
        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:
                    
                    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
        
        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:
            
            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):
        
        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)
        task.set_name(name)
        
        if task.done():
            self._on_task_done(task)
        else:
            self._tasks.add(task)
            task.add_done_callback(self._on_task_done)
        return task

    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():
            
            
            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:
            
            self._abort()
            self._parent_cancel_requested = True
            self._parent_task.cancel()