Accept async context managers for cleanup contexts (#11681)

CHANGES/11681.feature.rst

@@ -0,0 +1,10 @@
+Add support for using async context managers as cleanup contexts
+---------------------------------------------------------------
+
+Previously, application cleanup contexts required a single-yielding
+async generator. This change allows cleanup context callbacks to return
+contextlib.asynccontextmanager instances as well. The async context
+manager is adapted internally so it behaves like a single-yield async
+iterator: its ``__aenter__`` is executed at startup and ``__aexit__`` at
+cleanup. Existing async generator-based cleanup contexts remain
+supported.

CONTRIBUTORS.txt

@@ -289,6 +289,7 @@ Pahaz Blinov
 Panagiotis Kolokotronis
 Pankaj Pandey
 Parag Jain
+Parman Mohammadalizadeh
 Patrick Lee
 Pau Freixes
 Paul Colomiets

aiohttp/web_app.py

@@ -12,7 +12,15 @@
     Sequence,
 )
 from functools import lru_cache, partial, update_wrapper
-from typing import TYPE_CHECKING, Any, TypeVar, cast, final, overload
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncContextManager,
+    TypeVar,
+    cast,
+    final,
+    overload,
+)
 
 from aiosignal import Signal
 from frozenlist import FrozenList
@@ -140,7 +148,7 @@ def __init__(
 
     def __init_subclass__(cls: type["Application"]) -> None:
         raise TypeError(
-            f"Inheritance class {cls.__name__} from web.Application " "is forbidden"
+            f"Inheritance class {cls.__name__} from web.Application is forbidden"
         )
 
     # MutableMapping API
@@ -416,7 +424,12 @@ def exceptions(self) -> list[BaseException]:
 
 
 if TYPE_CHECKING:
-    _CleanupContextBase = FrozenList[Callable[[Application], AsyncIterator[None]]]
+    # cleanup contexts may be either async generators (async iterator)
+    # or async context managers (contextlib.asynccontextmanager). Accept
+    # both for type checking clarity.
+    _CleanupContextBase = FrozenList[
+        Callable[[Application], AsyncIterator[None] | AsyncContextManager[None]]
+    ]
 else:
     _CleanupContextBase = FrozenList
 
@@ -427,13 +440,38 @@ def __init__(self) -> None:
         self._exits: list[AsyncIterator[None]] = []
 
     async def _on_startup(self, app: Application) -> None:
+        """Run registered cleanup context callbacks at startup.
+
+        Each callback may return either an async iterator (an async
+        generator that yields exactly once) or an async context manager
+        (from contextlib.asynccontextmanager). If a context manager is
+        returned we wrap it in `_AsyncCMAsIterator` so it behaves like an
+        async iterator that yields once.
+
+        """
         for cb in self:
-            it = cb(app).__aiter__()
+            ctx = cb(app)
+            # If the callback returned an async context manager (has
+            # __aenter__ and __aexit__), wrap it so it behaves like an
+            # async iterator that yields once.
+            if hasattr(ctx, "__aenter__") and hasattr(ctx, "__aexit__"):
+                # Narrow type for mypy: ctx is an AsyncContextManager here.
+                it = _AsyncCMAsIterator(cast(AsyncContextManager[None], ctx))
+            else:
+                # Assume it's an async iterator / async generator
+                it = ctx.__aiter__()
+
             await it.__anext__()
             self._exits.append(it)
 
     async def _on_cleanup(self, app: Application) -> None:
-        errors = []
+        """Run cleanup for all registered contexts in reverse order.
+
+        Collects and re-raises exceptions similarly to previous
+        implementation: a single exception is propagated as-is, multiple
+        exceptions are wrapped into CleanupError.
+        """
+        errors: list[BaseException] = []
         for it in reversed(self._exits):
             try:
                 await it.__anext__()
@@ -448,3 +486,28 @@ async def _on_cleanup(self, app: Application) -> None:
                 raise errors[0]
             else:
                 raise CleanupError("Multiple errors on cleanup stage", errors)
+
+
+class _AsyncCMAsIterator:
+    """Wrap an async context manager instance to expose async iterator protocol used by CleanupContext.
+
+    The iterator will perform ``__aenter__`` on the first ``__anext__`` call
+    and ``__aexit__`` on the second one (which then raises StopAsyncIteration).
+    """
+
+    def __init__(self, cm: AsyncContextManager[None]) -> None:
+        self._cm = cm
+        self._entered = False
+
+    def __aiter__(self) -> "_AsyncCMAsIterator":
+        return self
+
+    async def __anext__(self) -> None:
+        if not self._entered:
+            # enter once and return control (equivalent to yielding once)
+            await self._cm.__aenter__()
+            self._entered = True
+            return None
+        # second call -> exit and stop iteration
+        await self._cm.__aexit__(None, None, None)
+        raise StopAsyncIteration

tests/test_cleanup_ctx.py

@@ -0,0 +1,45 @@
+from contextlib import asynccontextmanager
+from typing import AsyncIterator
+
+from aiohttp.web_app import Application
+
+
+async def test_cleanup_ctx_with_async_generator() -> None:
+    entered = []
+
+    async def ctx(app) -> AsyncIterator[None]:
+        entered.append("enter")
+        try:
+            yield
+        finally:
+            entered.append("exit")
+
+    app = Application()
+    app.cleanup_ctx.append(ctx)
+
+    await app._cleanup_ctx._on_startup(app)
+    assert entered == ["enter"]
+
+    await app._cleanup_ctx._on_cleanup(app)
+    assert entered == ["enter", "exit"]
+
+
+async def test_cleanup_ctx_with_asynccontextmanager() -> None:
+    entered = []
+
+    @asynccontextmanager
+    async def ctx(app) -> AsyncIterator[None]:
+        entered.append("enter")
+        try:
+            yield
+        finally:
+            entered.append("exit")
+
+    app = Application()
+    app.cleanup_ctx.append(ctx)
+
+    await app._cleanup_ctx._on_startup(app)
+    assert entered == ["enter"]
+
+    await app._cleanup_ctx._on_cleanup(app)
+    assert entered == ["enter", "exit"]