Class AsyncExitStack

Lib/contextlib.py:628–755

Async context manager for dynamic management of a stack of exit callbacks. For example: async with AsyncExitStack() as stack: connections = [await stack.enter_async_context(get_connection()) for i in range(5)] # All opened connections will aut

# Inspired by discussions on https://bugs.python.org/issue29302
class AsyncExitStack(_BaseExitStack, AbstractAsyncContextManager):
    """Async context manager for dynamic management of a stack of exit
    callbacks.

    For example:
        async with AsyncExitStack() as stack:
            connections = [await stack.enter_async_context(get_connection())
                for i in range(5)]
            # All opened connections will automatically be released at the
            # end of the async with statement, even if attempts to open a
            # connection later in the list raise an exception.
    """

    @staticmethod
    def _create_async_cb_wrapper(callback, /, *args, **kwds):
        async def _exit_wrapper(exc_type, exc, tb):
            await callback(*args, **kwds)
        return _exit_wrapper

    async def enter_async_context(self, cm):
        """Enters the supplied async context manager.

        If successful, also pushes its __aexit__ method as a callback and
        returns the result of the __aenter__ method.
        """
        _enter = _lookup_special(cm, '__aenter__', _sentinel)
        if _enter is _sentinel:
            cls = type(cm)
            raise TypeError(f"'{cls.__module__}.{cls.__qualname__}' object does "
                            f"not support the asynchronous context manager protocol")
        _exit = _lookup_special(cm, '__aexit__', _sentinel)
        if _exit is _sentinel:
            cls = type(cm)
            raise TypeError(f"'{cls.__module__}.{cls.__qualname__}' object does "
                            f"not support the asynchronous context manager protocol")
        result = await _enter()
        self._push_exit_callback(_exit, False)
        return result

    def push_async_exit(self, exit):
        """Registers a coroutine function with the standard __aexit__ method
        signature.

        Can suppress exceptions the same way __aexit__ method can.
        Also accepts any object with an __aexit__ method (registering a call
        to the method instead of the object itself).
        """
        exit_method = _lookup_special(exit, '__aexit__', exit)
        self._push_exit_callback(exit_method, False)
        return exit  # Allow use as a decorator

    def push_async_callback(self, callback, /, *args, **kwds):
        """Registers an arbitrary coroutine function and arguments.

        Cannot suppress exceptions.
        """
        _exit_wrapper = self._create_async_cb_wrapper(callback, *args, **kwds)