`BaseExceptionGroup` containing `CancelledError` catched by `asyncio.timeout()`

Async-SIG
1

Currently asyncio.timeout() cancels the inner block and raises up TimeoutError when it sees CancelledError in the __aexit__() handler and if it seems to be injected by it with a cancellation request count check.

When the inner block raises up an BaseExceptionGroup that contains the cancellation injected by timeout, it is not caught (because asyncio.timeouts.Timeout.__aexit__() just checks exc_typ is exceptions.CancelledError) and there is no way to distinguish cancellation by timeout and cancellation by other tasks.

The currently available asyncio.TaskGroup and other constructs seem to just raise up CancelledError upon timeouts, but this may not be the case. (Actually I discovered this while experimenting with extensions to Supervisor API… I’ll post a new discussion thread about its design. gather_safe() expermientation)

So, to seamlessly handle both naked exceptions and exceptions groups containing CancelledError raised up from the inner block in asyncio.timeout() and to replace CancelledError with TimeoutError, I had to change Timeout.__aexit__() as follows:

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_val: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> Optional[bool]:
        assert self._state in (_State.ENTERED, _State.EXPIRING)

        if self._timeout_handler is not None:
            self._timeout_handler.cancel()
            self._timeout_handler = None

        if self._state is _State.EXPIRING:
            self._state = _State.EXPIRED : Literal[_State.EXPIRED]

            contains_cancellation = False : Literal[False]
            if isinstance(exc_val, BaseExceptionGroup):
                matched, rest = exc_val.split(exceptions.CancelledError)
                contains_cancellation = (matched is not None) : bool
                if self._task.uncancel() <= self._cancelling and contains_cancellation:
                    if rest is not None:
                        raise BaseExceptionGroup("", [TimeoutError(), rest])
                    raise BaseExceptionGroup("", [TimeoutError()])
            else:
                if self._task.uncancel() <= self._cancelling and exc_type is exceptions.CancelledError:
                    # Since there are no new cancel requests, we're
                    # handling this.
                    raise TimeoutError from exc_val
        elif self._state is _State.ENTERED:
            self._state = _State.EXITED : Literal[_State.EXITED]

        return None

try:
    async with asyncio.timeout(1):
        await some_job_throwing_base_exception_group_when_cancelled()
except* asyncio.TimeoutError:
    print("timeout detected")
except* Exception:
    print("inner exception detected")

I think in general we need some way to reduce the boilerplates when re-raising exception groups and exceptions with replacement of specific exception types in them.

PEP-654’s use case examples mentions a scenario to combine inner exceptions and context manager’s exceptions, but it doesn’t mention about replacing cancellations like in the timeout context manager.

Questions

  • Would this be a special case for asyncio.timeout() or should every potential implementations of context managers take care of exception groups?
  • Would we need a generalization of above patterns?
    • e.g., Extend __exit__(), __aexit__() dunder methods to handle exception groups natively?
    • e.g., Add something like raise*? (related topic: python/exceptiongroups#7)

I’d like to hear other’s thoughts! (In particular, @yselivanov and @njs :pray: )

1 Like