codekingpro/portable-devtools

1# Adapted with permission from the EdgeDB project;2# license: PSFL.3 4 5__all__ = ("TaskGroup",)6 7from . import events8from . import exceptions9from . import futures10from . import tasks11 12 13class TaskGroup:14    """Asynchronous context manager for managing groups of tasks.15 16    Example use:17 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.")22 23    All tasks are awaited when the context manager exits.24 25    Any exceptions other than `asyncio.CancelledError` raised within26    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 = False31        self._exiting = False32        self._aborting = False33        self._loop = None34        self._parent_task = None35        self._parent_cancel_requested = False36        self._tasks = set()37        self._errors = []38        self._base_error = None39        self._on_completed_fut = None40 41    def __repr__(self):42        info = ['']43        if self._tasks:44            info.append(f'tasks={len(self._tasks)}')45        if self._errors:46            info.append(f'errors={len(self._errors)}')47        if self._aborting:48            info.append('cancelling')49        elif self._entered:50            info.append('entered')51 52        info_str = ' '.join(info)53        return f'<TaskGroup{info_str}>'54 55    async def __aenter__(self):56        if self._entered:57            raise RuntimeError(58                f"TaskGroup {self!r} has already been entered")59        if self._loop is None:60            self._loop = events.get_running_loop()61        self._parent_task = tasks.current_task(self._loop)62        if self._parent_task is None:63            raise RuntimeError(64                f'TaskGroup {self!r} cannot determine the parent task')65        self._entered = True66 67        return self68 69    async def __aexit__(self, et, exc, tb):70        tb = None71        try:72            return await self._aexit(et, exc)73        finally:74            # Exceptions are heavy objects that can have object75            # cycles (bad for GC); let's not keep a reference to76            # a bunch of them. It would be nicer to use a try/finally77            # in __aexit__ directly but that introduced some diff noise78            self._parent_task = None79            self._errors = None80            self._base_error = None81            exc = None82 83    async def _aexit(self, et, exc):84        self._exiting = True85 86        if (exc is not None and87                self._is_base_error(exc) and88                self._base_error is None):89            self._base_error = exc90 91        if et is not None and issubclass(et, exceptions.CancelledError):92            propagate_cancellation_error = exc93        else:94            propagate_cancellation_error = None95 96        if et is not None:97            if not self._aborting:98                # Our parent task is being cancelled:99                #100                #    async with TaskGroup() as g:101                #        g.create_task(...)102                #        await ...  # <- CancelledError103                #104                # or there's an exception in "async with":105                #106                #    async with TaskGroup() as g:107                #        g.create_task(...)108                #        1 / 0109                #110                self._abort()111 112        # We use while-loop here because "self._on_completed_fut"113        # can be cancelled multiple times if our parent task114        # is being cancelled repeatedly (or even once, when115        # our own cancellation is already in progress)116        while self._tasks:117            if self._on_completed_fut is None:118                self._on_completed_fut = self._loop.create_future()119 120            try:121                await self._on_completed_fut122            except exceptions.CancelledError as ex:123                if not self._aborting:124                    # Our parent task is being cancelled:125                    #126                    #    async def wrapper():127                    #        async with TaskGroup() as g:128                    #            g.create_task(foo)129                    #130                    # "wrapper" is being cancelled while "foo" is131                    # still running.132                    propagate_cancellation_error = ex133                    self._abort()134 135            self._on_completed_fut = None136 137        assert not self._tasks138 139        if self._base_error is not None:140            try:141                raise self._base_error142            finally:143                exc = None144 145        if self._parent_cancel_requested:146            # If this flag is set we *must* call uncancel().147            if self._parent_task.uncancel() == 0:148                # If there are no pending cancellations left,149                # don't propagate CancelledError.150                propagate_cancellation_error = None151 152        # Propagate CancelledError if there is one, except if there153        # are other errors -- those have priority.154        try:155            if propagate_cancellation_error is not None and not self._errors:156                try:157                    raise propagate_cancellation_error158                finally:159                    exc = None160        finally:161            propagate_cancellation_error = None162 163        if et is not None and not issubclass(et, exceptions.CancelledError):164            self._errors.append(exc)165 166        if self._errors:167            # If the parent task is being cancelled from the outside168            # of the taskgroup, un-cancel and re-cancel the parent task,169            # which will keep the cancel count stable.170            if self._parent_task.cancelling():171                self._parent_task.uncancel()172                self._parent_task.cancel()173            try:174                raise BaseExceptionGroup(175                    'unhandled errors in a TaskGroup',176                    self._errors,177                ) from None178            finally:179                exc = None180 181 182    def create_task(self, coro, **kwargs):183        """Create a new task in this group and return it.184 185        Similar to `asyncio.create_task`.186        """187        if not self._entered:188            coro.close()189            raise RuntimeError(f"TaskGroup {self!r} has not been entered")190        if self._exiting and not self._tasks:191            coro.close()192            raise RuntimeError(f"TaskGroup {self!r} is finished")193        if self._aborting:194            coro.close()195            raise RuntimeError(f"TaskGroup {self!r} is shutting down")196        task = self._loop.create_task(coro, **kwargs)197 198        futures.future_add_to_awaited_by(task, self._parent_task)199 200        # Always schedule the done callback even if the task is201        # already done (e.g. if the coro was able to complete eagerly),202        # otherwise if the task completes with an exception then it will cancel203        # the current task too early. gh-128550, gh-128588204        self._tasks.add(task)205        task.add_done_callback(self._on_task_done)206        try:207            return task208        finally:209            # gh-128552: prevent a refcycle of210            # task.exception().__traceback__->TaskGroup.create_task->task211            del task212 213    # Since Python 3.8 Tasks propagate all exceptions correctly,214    # except for KeyboardInterrupt and SystemExit which are215    # still considered special.216 217    def _is_base_error(self, exc: BaseException) -> bool:218        assert isinstance(exc, BaseException)219        return isinstance(exc, (SystemExit, KeyboardInterrupt))220 221    def _abort(self):222        self._aborting = True223 224        for t in self._tasks:225            if not t.done():226                t.cancel()227 228    def _on_task_done(self, task):229        self._tasks.discard(task)230 231        futures.future_discard_from_awaited_by(task, self._parent_task)232 233        if self._on_completed_fut is not None and not self._tasks:234            if not self._on_completed_fut.done():235                self._on_completed_fut.set_result(True)236 237        if task.cancelled():238            return239 240        exc = task.exception()241        if exc is None:242            return243 244        self._errors.append(exc)245        if self._is_base_error(exc) and self._base_error is None:246            self._base_error = exc247 248        if self._parent_task.done():249            # Not sure if this case is possible, but we want to handle250            # it anyways.251            self._loop.call_exception_handler({252                'message': f'Task {task!r} has errored out but its parent '253                           f'task {self._parent_task} is already completed',254                'exception': exc,255                'task': task,256            })257            return258 259        if not self._aborting and not self._parent_cancel_requested:260            # If parent task *is not* being cancelled, it means that we want261            # to manually cancel it to abort whatever is being run right now262            # in the TaskGroup.  But we want to mark parent task as263            # "not cancelled" later in __aexit__.  Example situation that264            # we need to handle:265            #266            #    async def foo():267            #        try:268            #            async with TaskGroup() as g:269            #                g.create_task(crash_soon())270            #                await something  # <- this needs to be canceled271            #                                 #    by the TaskGroup, e.g.272            #                                 #    foo() needs to be cancelled273            #        except Exception:274            #            # Ignore any exceptions raised in the TaskGroup275            #            pass276            #        await something_else     # this line has to be called277            #                                 # after TaskGroup is finished.278            self._abort()279            self._parent_cancel_requested = True280            self._parent_task.cancel()281