Added context manager mix-in classes

docs/api.rst

@@ -95,6 +95,15 @@ Temporary files and directories
 .. autoclass:: anyio.SpooledTemporaryFile
 .. autoclass:: anyio.TemporaryDirectory
 
+Context manager mix-in classes
+------------------------------
+
+.. autoclass:: anyio.ContextManagerMixin
+   :special-members: __contextmanager__
+
+.. autoclass:: anyio.AsyncContextManagerMixin
+   :special-members: __asynccontextmanager__
+
 Streams and stream wrappers
 ---------------------------
 

docs/cancellation.rst

@@ -196,8 +196,11 @@ Depending on how they are used, this pattern is, however, *usually* safe to use
 asynchronous context managers, so long as you make sure that the same host task keeps
 running throughout the entire enclosed code block::
 
+    from contextlib import asynccontextmanager
+
+
     # Okay in most cases!
-    @async_context_manager
+    @asynccontextmanager
     async def some_context_manager():
         async with create_task_group() as tg:
             tg.start_soon(foo)
@@ -209,24 +212,20 @@ start to end in the same task, making it possible to have task groups or cancel
 safely straddle the ``yield``.
 
 When you're implementing the async context manager protocol manually and your async
-context manager needs to use other context managers, you may find it necessary to call
-their ``__aenter__()`` and ``__aexit__()`` directly. In such cases, it is absolutely
-vital to ensure that their ``__aexit__()`` methods are called in the exact reverse order
-of the ``__aenter__()`` calls. To this end, you may find the
-:class:`~contextlib.AsyncExitStack` class very useful::
+context manager needs to use other context managers, you may find it convenient to use
+:class:`AsyncContextManagerMixin` in order to avoid cumbersome code that calls
+``__aenter__()`` and ``__aexit__()`` directly::
 
-    from contextlib import AsyncExitStack
+    from __future__ import annotations
 
-    from anyio import create_task_group
+    from collections.abc import AsyncGenerator
 
+    from anyio import AsyncContextManagerMixin, create_task_group
 
-    class MyAsyncContextManager:
-        async def __aenter__(self):
-            self._exitstack = AsyncExitStack()
-            await self._exitstack.__aenter__()
-            self._task_group = await self._exitstack.enter_async_context(
-                create_task_group()
-            )
 
-        async def __aexit__(self, exc_type, exc_val, exc_tb):
-            return await self._exitstack.__aexit__(exc_type, exc_val, exc_tb)
+    # AsyncContextManagerMixin is parametrized this way because it returns 'self'
+    class MyAsyncContextManager(AsyncContextManagerMixin["MyAsyncContextManager"]):
+        async def __asynccontextmanager__(self) -> AsyncGenerator[MyAsyncContextManager]:
+            async with create_task_group() as tg:
+                ...  # launch tasks
+                yield self

docs/versionhistory.rst

@@ -3,6 +3,12 @@ Version history
 
 This library adheres to `Semantic Versioning 2.0 <http://semver.org/>`_.
 
+**UNRELEASED**
+
+- Added context manager mix-in classes (``anyio.ContextManagerMixin`` and
+  ``anyio.AsyncContextManagerMixin``) to help write classes that embed other context
+  managers (particularly cancel scopes or task groups)
+
 **4.9.0**
 
 - Added async support for temporary file handling

pyproject.toml

@@ -127,6 +127,10 @@ relative_files = true
 
 [tool.coverage.report]
 show_missing = true
+exclude_also = [
+    "if TYPE_CHECKING:",
+    "@(abc\\.)?abstractmethod",
+]
 
 [tool.tox]
 env_list = ["pre-commit", "py39", "py310", "py311", "py312", "py313", "py314", "pypy3"]

src/anyio/__init__.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from ._core._contextmanagers import AsyncContextManagerMixin as AsyncContextManagerMixin
+from ._core._contextmanagers import ContextManagerMixin as ContextManagerMixin
 from ._core._eventloop import current_time as current_time
 from ._core._eventloop import get_all_backends as get_all_backends
 from ._core._eventloop import get_cancelled_exc_class as get_cancelled_exc_class

src/anyio/_core/_contextmanagers.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+from abc import abstractmethod
+from collections.abc import AsyncGenerator, Generator
+from inspect import isasyncgen, iscoroutine
+from types import TracebackType
+from typing import Generic, TypeVar, final
+
+_T_co = TypeVar("_T_co", covariant=True)
+
+
+class ContextManagerMixin(Generic[_T_co]):
+    """
+    Mixin class providing context manager functionality via a generator-based
+    implementation.
+
+    This class allows you to implement a context manager via :meth:`__contextmanager__`
+    which should return a generator. The mechanics are meant to mirror those of
+    :func:`@contextmanager <contextlib.contextmanager>`.
+    """
+
+    @final
+    def __enter__(self) -> _T_co:
+        gen = self.__contextmanager__()
+        if not isinstance(gen, Generator):
+            raise TypeError(
+                f"__contextmanager__() did not return a generator object, "
+                f"but {gen.__class__!r}"
+            )
+
+        try:
+            value = gen.send(None)
+        except StopIteration:
+            raise RuntimeError(
+                "the __contextmanager__() generator returned without yielding a value"
+            ) from None
+
+        self.__cm = gen
+        return value
+
+    @final
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool | None:
+        # Prevent circular references
+        cm = self.__cm
+        del self.__cm
+
+        if exc_val is not None:
+            try:
+                cm.throw(exc_val)
+            except StopIteration:
+                return True
+        else:
+            try:
+                cm.send(None)
+            except StopIteration:
+                return None
+
+        cm.close()
+        raise RuntimeError("the __contextmanager__() generator didn't stop")
+
+    @abstractmethod
+    def __contextmanager__(self) -> Generator[_T_co, None, None]:
+        """
+        Implement your context manager logic here, as you would with
+        :func:`@contextmanager <contextlib.contextmanager>`.
+
+        Any code up to the ``yield`` will be run in ``__enter__()``, and any code after
+        it is run in ``__exit__()``.
+
+        .. note:: If an exception is raised in the context block, it is reraised from
+            the ``yield``, just like with
+            :func:`@contextmanager <contextlib.contextmanager>`.
+
+        :return: a generator that yields exactly once
+        """
+
+
+class AsyncContextManagerMixin(Generic[_T_co]):
+    """
+    Mixin class providing async context manager functionality via a generator-based
+    implementation.
+
+    This class allows you to implement a context manager via
+    :meth:`__asynccontextmanager__`. The mechanics are meant to mirror those of
+    :func:`@asynccontextmanager <contextlib.asynccontextmanager>`.
+    """
+
+    @final
+    async def __aenter__(self) -> _T_co:
+        gen = self.__asynccontextmanager__()
+        if not isasyncgen(gen):
+            if iscoroutine(gen):
+                gen.close()
+                raise TypeError(
+                    "__asynccontextmanager__() returned a coroutine object instead of "
+                    "an async generator. Did you forget to add 'yield'?"
+                )
+
+            raise TypeError(
+                f"__asynccontextmanager__() did not return an async generator object, "
+                f"but {gen.__class__!r}"
+            )
+
+        try:
+            value = await gen.asend(None)
+        except StopAsyncIteration:
+            raise RuntimeError(
+                "the __asynccontextmanager__() generator returned without yielding a "
+                "value"
+            ) from None
+
+        self.__cm = gen
+        return value
+
+    @final
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool | None:
+        # Prevent circular references
+        cm = self.__cm
+        del self.__cm
+
+        if exc_val is not None:
+            try:
+                await cm.athrow(exc_val)
+            except StopAsyncIteration:
+                return True
+        else:
+            try:
+                await cm.asend(None)
+            except StopAsyncIteration:
+                return None
+
+        await cm.aclose()
+        raise RuntimeError("the __asynccontextmanager__() generator didn't stop")
+
+    @abstractmethod
+    def __asynccontextmanager__(self) -> AsyncGenerator[_T_co, None]:
+        """
+        Implement your async context manager logic here, as you would with
+        :func:`@asynccontextmanager <contextlib.asynccontextmanager>`.
+
+        Any code up to the ``yield`` will be run in ``__aenter__()``, and any code after
+        it is run in ``__aexit__()``.
+
+        .. note:: If an exception is raised in the context block, it is reraised from
+            the ``yield``, just like with
+            :func:`@asynccontextmanager <contextlib.asynccontextmanager>`.
+
+        :return: an async generator that yields exactly once
+        """