Repository: agronholm/anyio
Pull request #912: Fix contextmanager mixin typing
Author: login: tapetersen
Created: 2025-04-16T13:18:21Z | Merged: 2025-04-18T11:34:07Z

<!-- Thank you for your contribution! -->
## Changes

Fixes #. <!-- Provide issue number if exists -->

<!-- Please give a short brief about these changes. -->

## Checklist

If this is a user-facing code change, like a bugfix or a new feature, please ensure that
you've fulfilled the following conditions (where applicable):

- [ ] You've added tests (in `tests/`) added which would fail without your patch
- [ ] You've updated the documentation (in `docs/`, in case of behavior changes or new
features)
- [ ] You've added a new changelog entry (in `docs/versionhistory.rst`).

If this is a trivial change, like a typo fix or a code reformatting, then you can ignore
these instructions.

### Updating the changelog

If there are no entries after the last release, use `**UNRELEASED**` as the version.
If, say, your patch fixes issue <span>#</span>123, the entry should look like this:

```
- Fix big bad boo-boo in task groups
  (`#123 <https://github.com/agronholm/anyio/issues/123>`_; PR by @yourgithubaccount)
```

If there's no issue linked, just link to your pull request instead by updating the
changelog after you've created the PR.



Diff:
diff --git a/src/anyio/_core/_contextmanagers.py b/src/anyio/_core/_contextmanagers.py
index 25ec05a2b..0c3ef4120 100644
--- a/src/anyio/_core/_contextmanagers.py
+++ b/src/anyio/_core/_contextmanagers.py
@@ -4,12 +4,23 @@
 from contextlib import AbstractAsyncContextManager, AbstractContextManager
 from inspect import isasyncgen, iscoroutine, isgenerator
 from types import TracebackType
-from typing import Generic, TypeVar, final
+from typing import Protocol, TypeVar, cast, final
 
 _T_co = TypeVar("_T_co", covariant=True)
+_ExitT_co = TypeVar("_ExitT_co", covariant=True, bound="bool | None")
 
 
-class ContextManagerMixin(Generic[_T_co]):
+class _SupportsCtxMgr(Protocol[_T_co, _ExitT_co]):
+    def __contextmanager__(self) -> AbstractContextManager[_T_co, _ExitT_co]: ...
+
+
+class _SupportsAsyncCtxMgr(Protocol[_T_co, _ExitT_co]):
+    def __asynccontextmanager__(
+        self,
+    ) -> AbstractAsyncContextManager[_T_co, _ExitT_co]: ...
+
+
+class ContextManagerMixin:
     """
     Mixin class providing context manager functionality via a generator-based
     implementation.
@@ -22,10 +33,12 @@ class ContextManagerMixin(Generic[_T_co]):
         that once you enter it, you can't re-enter before first exiting it.
     """
 
-    __cm: AbstractContextManager[_T_co] | None = None
+    __cm: AbstractContextManager[object, bool | None] | None = None
 
     @final
-    def __enter__(self) -> _T_co:
+    def __enter__(self: _SupportsCtxMgr[_T_co, bool | None]) -> _T_co:
+        # Needed for mypy to assume self still has the __cm member
+        assert isinstance(self, ContextManagerMixin)
         if self.__cm is not None:
             raise RuntimeError(
                 f"this {self.__class__.__qualname__} has already been entered"
@@ -58,11 +71,13 @@ def __enter__(self) -> _T_co:
 
     @final
     def __exit__(
-        self,
+        self: _SupportsCtxMgr[object, _ExitT_co],
         exc_type: type[BaseException] | None,
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
-    ) -> bool | None:
+    ) -> _ExitT_co:
+        # Needed for mypy to assume self still has the __cm member
+        assert isinstance(self, ContextManagerMixin)
         if self.__cm is None:
             raise RuntimeError(
                 f"this {self.__class__.__qualname__} has not been entered yet"
@@ -72,10 +87,10 @@ def __exit__(
         cm = self.__cm
         del self.__cm
 
-        return cm.__exit__(exc_type, exc_val, exc_tb)
+        return cast(_ExitT_co, cm.__exit__(exc_type, exc_val, exc_tb))
 
     @abstractmethod
-    def __contextmanager__(self) -> AbstractContextManager[_T_co]:
+    def __contextmanager__(self) -> AbstractContextManager[object, bool | None]:
         """
         Implement your context manager logic here.
 
@@ -89,7 +104,7 @@ def __contextmanager__(self) -> AbstractContextManager[_T_co]:
         """
 
 
-class AsyncContextManagerMixin(Generic[_T_co]):
+class AsyncContextManagerMixin:
     """
     Mixin class providing async context manager functionality via a generator-based
     implementation.
@@ -102,10 +117,12 @@ class AsyncContextManagerMixin(Generic[_T_co]):
         that once you enter it, you can't re-enter before first exiting it.
     """
 
-    __cm: AbstractAsyncContextManager[_T_co] | None = None
+    __cm: AbstractAsyncContextManager[object, bool | None] | None = None
 
     @final
-    async def __aenter__(self) -> _T_co:
+    async def __aenter__(self: _SupportsAsyncCtxMgr[_T_co, bool | None]) -> _T_co:
+        # Needed for mypy to assume self still has the __cm member
+        assert isinstance(self, AsyncContextManagerMixin)
         if self.__cm is not None:
             raise RuntimeError(
                 f"this {self.__class__.__qualname__} has already been entered"
@@ -145,11 +162,12 @@ async def __aenter__(self) -> _T_co:
 
     @final
     async def __aexit__(
-        self,
+        self: _SupportsAsyncCtxMgr[object, _ExitT_co],
         exc_type: type[BaseException] | None,
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
-    ) -> bool | None:
+    ) -> _ExitT_co:
+        assert isinstance(self, AsyncContextManagerMixin)
         if self.__cm is None:
             raise RuntimeError(
                 f"this {self.__class__.__qualname__} has not been entered yet"
@@ -159,10 +177,12 @@ async def __aexit__(
         cm = self.__cm
         del self.__cm
 
-        return await cm.__aexit__(exc_type, exc_val, exc_tb)
+        return cast(_ExitT_co, await cm.__aexit__(exc_type, exc_val, exc_tb))
 
     @abstractmethod
-    def __asynccontextmanager__(self) -> AbstractAsyncContextManager[_T_co]:
+    def __asynccontextmanager__(
+        self,
+    ) -> AbstractAsyncContextManager[object, bool | None]:
         """
         Implement your async context manager logic here.
 
diff --git a/tests/test_contextmanagers.py b/tests/test_contextmanagers.py
index 5e7191177..1ce36b989 100644
--- a/tests/test_contextmanagers.py
+++ b/tests/test_contextmanagers.py
@@ -22,7 +22,7 @@
 pytestmark = pytest.mark.anyio
 
 
-class DummyContextManager(ContextManagerMixin["DummyContextManager"]):
+class DummyContextManager(ContextManagerMixin):
     def __init__(self, handle_exc: bool = False) -> None:
         self.started = False
         self.finished = False
@@ -40,7 +40,7 @@ def __contextmanager__(self) -> Generator[Self]:
         self.finished = True
 
 
-class DummyAsyncContextManager(AsyncContextManagerMixin["DummyAsyncContextManager"]):
+class DummyAsyncContextManager(AsyncContextManagerMixin):
     def __init__(self, handle_exc: bool = False) -> None:
         self.started = False
         self.finished = False
@@ -79,7 +79,7 @@ def test_exception(self, handle_exc: bool) -> None:
         assert cm.finished == handle_exc
 
     def test_return_bad_type(self) -> None:
-        class BadContextManager(ContextManagerMixin[None]):
+        class BadContextManager(ContextManagerMixin):
             def __contextmanager__(self) -> AbstractContextManager[None]:
                 return None  # type: ignore[return-value]
 
@@ -88,7 +88,7 @@ def __contextmanager__(self) -> AbstractContextManager[None]:
                 pass
 
     def test_return_generator(self) -> None:
-        class BadContextManager(ContextManagerMixin["BadContextManager"]):
+        class BadContextManager(ContextManagerMixin):
             def __contextmanager__(self):  # type: ignore[no-untyped-def]
                 yield self
 
@@ -97,7 +97,7 @@ def __contextmanager__(self):  # type: ignore[no-untyped-def]
                 pass
 
     def test_return_self(self) -> None:
-        class BadContextManager(ContextManagerMixin["BadContextManager"]):
+        class BadContextManager(ContextManagerMixin):
             def __contextmanager__(self):  # type: ignore[no-untyped-def]
                 return self
 
@@ -173,7 +173,7 @@ async def test_exception(self, handle_exc: bool) -> None:
         assert cm.finished == handle_exc
 
     async def test_return_bad_type(self) -> None:
-        class BadContextManager(AsyncContextManagerMixin[None]):
+        class BadContextManager(AsyncContextManagerMixin):
             def __asynccontextmanager__(self):  # type: ignore[no-untyped-def]
                 return None
 
@@ -182,7 +182,7 @@ def __asynccontextmanager__(self):  # type: ignore[no-untyped-def]
                 pass
 
     async def test_return_async_generator(self) -> None:
-        class BadContextManager(AsyncContextManagerMixin["BadContextManager"]):
+        class BadContextManager(AsyncContextManagerMixin):
             async def __asynccontextmanager__(self):  # type: ignore[no-untyped-def]
                 yield self
 
@@ -191,7 +191,7 @@ async def __asynccontextmanager__(self):  # type: ignore[no-untyped-def]
                 pass
 
     async def test_return_self(self) -> None:
-        class BadContextManager(AsyncContextManagerMixin["BadContextManager"]):
+        class BadContextManager(AsyncContextManagerMixin):
             def __asynccontextmanager__(self):  # type: ignore[no-untyped-def]
                 return self
 
@@ -200,7 +200,7 @@ def __asynccontextmanager__(self):  # type: ignore[no-untyped-def]
                 pass
 
     async def test_return_coroutine(self) -> None:
-        class BadContextManager(AsyncContextManagerMixin["BadContextManager"]):
+        class BadContextManager(AsyncContextManagerMixin):
             async def __asynccontextmanager__(self):  # type: ignore[no-untyped-def]
                 return self
 


Reviews:
user: login: tapetersen
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-16T18:39:59Z
user: login: tapetersen
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-16T18:42:54Z
user: login: agronholm
body: I can't find fault in this, so I'm merging it.
state: APPROVED
author_association: OWNER
submitted_at: 2025-04-18T11:33:55Z

Review comments:
diff_hunk: @@ -4,12 +4,20 @@
 from contextlib import AbstractAsyncContextManager, AbstractContextManager
 from inspect import isasyncgen, iscoroutine, isgenerator
 from types import TracebackType
-from typing import Generic, TypeVar, final
+from typing import Protocol, TypeVar, final
 
 _T_co = TypeVar("_T_co", covariant=True)
 
 
-class ContextManagerMixin(Generic[_T_co]):
+class _SupportsCtxMgr(Protocol[_T_co]):
+    def __contextmanager__(self) -> AbstractContextManager[_T_co]: ...
+
+
+class _SupportsAsyncCtxMgr(Protocol[_T_co]):
+    def __asynccontextmanager__(self) -> AbstractAsyncContextManager[_T_co]: ...
+
+
+class ContextManagerMixin:
path: src/anyio/_core/_contextmanagers.py
user: login: tapetersen
body: Removing the typevar from the class level here is not necessary but avoids having to declare another one to use in the method and since the mixin class doesn't really depend on the value type it's not needed.
created_at: 2025-04-16T18:39:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -102,10 +112,12 @@ class AsyncContextManagerMixin(Generic[_T_co]):
         that once you enter it, you can't re-enter before first exiting it.
     """
 
-    __cm: AbstractAsyncContextManager[_T_co] | None = None
+    __cm: AbstractAsyncContextManager[object] | None = None
 
     @final
-    async def __aenter__(self) -> _T_co:
+    async def __aenter__(self: _SupportsAsyncCtxMgr[_T_co]) -> _T_co:
+        # Needed for mypy to assume self still has the __cm member
+        assert isinstance(self, AsyncContextManagerMixin)
path: src/anyio/_core/_contextmanagers.py
user: login: tapetersen
body: pyright treats the self typing as an implicit intersection with the containing class but mypy doesn't (even if it does other magic inference fine).
created_at: 2025-04-16T18:42:54Z
author_association: CONTRIBUTOR

PR comments:
user: login: agronholm
created_at: 2025-04-16T18:03:15Z
body: Just so you know, I really don't think I will merge this as I have an alternate solution on the way.
author_association: OWNER
user: login: tapetersen
created_at: 2025-04-16T18:16:16Z
body: Sure. I'll leave it in that case. This does seem to work and let's the inference work without any explicit parametrization at all though. 
author_association: CONTRIBUTOR
user: login: agronholm
created_at: 2025-04-16T18:17:04Z
body: Oh, you improved it? In that case I'll take a closer look.
author_association: OWNER
user: login: tapetersen
created_at: 2025-04-16T18:36:15Z
body: Yeah, as I wrote in the other thread I didn't really think the old was an improvement but I have used the trick of typing self before for other cases and thought to give it a try.

I'm somewhat surprised it worked as well as it did in mypy since I've mainly used these tricks with pyright and expected it to be possibly problematic.

The lack of type-intersection at least for protocols is really hurting here though... (it would allow this to be implemented with a decorator instead of a mixin but there is no way to specify the return type of that correctly unfortunately)

author_association: CONTRIBUTOR
user: login: agronholm
created_at: 2025-04-16T21:36:50Z
body: I'm leaning towards accepting this, as it seems to work exactly the way I wanted! Thanks! I'll give this a more thorough review within the next couple days.
author_association: OWNER