Exempt async CM methods from ASYNC910/911 when partner checkpoints

ASYNC910 and ASYNC911 no longer require every `__aenter__`/`__aexit__` to
contain a checkpoint. Per Trio's documentation, an async context manager
only needs one of entry/exit to act as a checkpoint.

When a class defines both methods, the one without an `await` is exempt
if its partner contains one. When a class defines only one of the two,
the partner is charitably assumed to be inherited from a base class and
to contain a checkpoint, so the defined method is also exempt.

Closes #441

https://claude.ai/code/session_014jAydKywq31Ew4fVYGJdiG

Author: claude

docs/changelog.rst

@@ -7,6 +7,7 @@ Changelog
 Unreleased
 ==========
 - Autofix for :ref:`ASYNC910 <async910>` / :ref:`ASYNC911 <async911>` no longer inserts checkpoints inside ``except`` clauses (which would trigger :ref:`ASYNC120 <async120>`); instead the checkpoint is added at the top of the function or of the enclosing loop. `(issue #403) <https://github.com/python-trio/flake8-async/issues/403>`_
+- :ref:`ASYNC910 <async910>` and :ref:`ASYNC911 <async911>` now accept ``__aenter__`` / ``__aexit__`` methods when the partner method provides the checkpoint, or when only one of the two is defined on a class that inherits from another class (charitably assuming the partner is inherited and contains a checkpoint). `(issue #441) <https://github.com/python-trio/flake8-async/issues/441>`_
 
 25.7.1
 ======

flake8_async/visitors/visitor91x.py

@@ -465,6 +465,17 @@ def __init__(self, *args: Any, **kwargs: Any):
         # used to transfer new body between visit_FunctionDef and leave_FunctionDef
         self.new_body: cst.BaseSuite | None = None
 
+        # Tracks whether the current scope is a class body and, if so, which of
+        # `__aenter__`/`__aexit__` are directly defined on it (values: True if
+        # that method contains an `await`, False otherwise, missing key if not
+        # defined). Used to exempt async context manager methods from
+        # ASYNC910/911 when their partner method provides the checkpoint, or
+        # when the partner is assumed inherited (not defined on this class).
+        self.async_cm_class: dict[str, bool] | None = None
+        # Set on entry to an exempt `__aenter__`/`__aexit__` so that
+        # `error_91x` skips emitting ASYNC910/911.
+        self.exempt_async_cm_method = False
+
     def should_autofix(self, node: cst.CSTNode, code: str | None = None) -> bool:
         if code is None:
             code = "ASYNC911" if self.has_yield else "ASYNC910"
@@ -532,6 +543,44 @@ def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
                         self.suppress_imported_as.append("suppress")
                     return
 
+    # Async context manager methods may legitimately skip checkpointing if the
+    # partner method provides the checkpoint, or if the partner is inherited
+    # from a base class (which we charitably assume contains a checkpoint).
+    # See https://github.com/python-trio/flake8-async/issues/441.
+    def visit_ClassDef(self, node: cst.ClassDef) -> None:
+        self.save_state(node, "async_cm_class")
+        defined: dict[str, bool] = {}
+        if isinstance(node.body, cst.IndentedBlock):
+            for stmt in node.body.body:
+                if (
+                    isinstance(stmt, cst.FunctionDef)
+                    and stmt.asynchronous is not None
+                    and stmt.name.value in ("__aenter__", "__aexit__")
+                ):
+                    defined[stmt.name.value] = bool(m.findall(stmt, m.Await()))
+        self.async_cm_class = defined
+
+    def leave_ClassDef(
+        self, original_node: cst.ClassDef, updated_node: cst.ClassDef
+    ) -> cst.ClassDef:
+        self.restore_state(original_node)
+        return updated_node
+
+    def _is_exempt_async_cm_method(self, node: cst.FunctionDef) -> bool:
+        if self.async_cm_class is None:
+            return False
+        name = node.name.value
+        if name not in ("__aenter__", "__aexit__"):
+            return False
+        if name not in self.async_cm_class:
+            return False
+        partner = "__aexit__" if name == "__aenter__" else "__aenter__"
+        # Partner not defined in this class -> assume inherited with checkpoint.
+        if partner not in self.async_cm_class:
+            return True
+        # Partner defined and (charitably) contains a checkpoint.
+        return self.async_cm_class[partner]
+
     def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
         # `await` in default values happen in parent scope
         # we also know we don't ever modify parameters so we can ignore the return value
@@ -543,6 +592,8 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
         if func_has_decorator(node, "overload", "fixture") or func_empty_body(node):
             return False  # subnodes can be ignored
 
+        is_exempt_cm = self._is_exempt_async_cm_method(node)
+
         self.save_state(
             node,
             "has_yield",
@@ -557,6 +608,8 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
             "suppress_imported_as",  # a copy is saved, but state is not reset
             "except_depth",
             "add_checkpoint_at_function_start",
+            "async_cm_class",
+            "exempt_async_cm_method",
             copy=True,
         )
         self.uncheckpointed_statements = set()
@@ -567,6 +620,9 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
         self.taskgroup_has_start_soon = {}
         self.except_depth = 0
         self.add_checkpoint_at_function_start = False
+        # Class-level context does not apply to nested scopes.
+        self.async_cm_class = None
+        self.exempt_async_cm_method = is_exempt_cm
 
         self.async_function = (
             node.asynchronous is not None
@@ -747,6 +803,12 @@ def error_91x(
     ) -> bool:
         assert not isinstance(statement, ArtificialStatement), statement
 
+        # Exempt `__aenter__`/`__aexit__` when the partner method contains a
+        # checkpoint, or when the partner is missing and charitably assumed
+        # inherited.
+        if self.exempt_async_cm_method:
+            return False
+
         if isinstance(node, cst.FunctionDef):
             msg = "exit"
         else:

tests/autofix_files/async910.py

@@ -636,3 +636,83 @@ async def foo_nested_empty_async():
     async def bar(): ...
 
     await foo()
+
+
+# Issue #441: async context manager methods may legitimately skip checkpointing
+# if the partner method provides the checkpoint, or if the partner is inherited.
+class CtxWithSetup:  # safe: __aenter__ checkpoints, __aexit__ can be fast
+    async def __aenter__(self):
+        await foo()
+
+    async def __aexit__(self, exc_type, exc, tb):
+        print("fast exit")
+
+
+class CtxWithTeardown:  # safe: __aexit__ checkpoints, __aenter__ can be fast
+    async def __aenter__(self):
+        print("fast setup")
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await foo()
+
+
+class CtxWithBothCheckpoint:  # safe: both checkpoint
+    async def __aenter__(self):
+        await foo()
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await foo()
+
+
+# fmt: off
+class CtxNeitherCheckpoint:
+    async def __aenter__(self):  # error: 4, "exit", Stmt("function definition", line)
+        print("setup")
+        await trio.lowlevel.checkpoint()
+
+    async def __aexit__(self, *a):  # error: 4, "exit", Stmt("function definition", line)
+        print("teardown")
+        await trio.lowlevel.checkpoint()
+# fmt: on
+
+
+# Only one method defined: charitably assume the other is inherited with a checkpoint.
+class CtxOnlyAenter:  # safe: __aexit__ assumed inherited with checkpoint
+    async def __aenter__(self):
+        print("setup")
+
+
+class CtxOnlyAexit:  # safe: __aenter__ assumed inherited with checkpoint
+    async def __aexit__(self, *a):
+        print("teardown")
+
+
+class CtxOnlyAenterWithCheckpoint:  # safe
+    async def __aenter__(self):
+        await foo()
+
+
+class CtxOnlyAexitWithCheckpoint:  # safe
+    async def __aexit__(self, *a):
+        await foo()
+
+
+# a nested function named `__aenter__` inside another function is not a method
+def not_a_class():
+    async def __aenter__(self):  # error: 4, "exit", Stmt("function definition", line)
+        print("setup")
+        await trio.lowlevel.checkpoint()
+
+
+# class nested inside a function still gets the exemption
+def factory():
+    class NestedCtx:  # safe
+        async def __aenter__(self):
+            print("setup")
+
+
+# nested class; outer class has nothing relevant
+class Outer:
+    class Inner:  # safe: charitable inheritance for __aexit__
+        async def __aenter__(self):
+            print("setup")

tests/autofix_files/async910.py.diff

@@ -223,3 +223,23 @@
 
 
  async def foo_nested_empty_async():
+@@ x,9 x,11 @@
+ class CtxNeitherCheckpoint:
+     async def __aenter__(self):  # error: 4, "exit", Stmt("function definition", line)
+         print("setup")
++        await trio.lowlevel.checkpoint()
+
+     async def __aexit__(self, *a):  # error: 4, "exit", Stmt("function definition", line)
+         print("teardown")
++        await trio.lowlevel.checkpoint()
+ # fmt: on
+
+
+@@ x,6 x,7 @@
+ def not_a_class():
+     async def __aenter__(self):  # error: 4, "exit", Stmt("function definition", line)
+         print("setup")
++        await trio.lowlevel.checkpoint()
+
+
+ # class nested inside a function still gets the exemption

tests/eval_files/async910.py

@@ -606,3 +606,80 @@ async def foo_nested_empty_async():
     async def bar(): ...
 
     await foo()
+
+
+# Issue #441: async context manager methods may legitimately skip checkpointing
+# if the partner method provides the checkpoint, or if the partner is inherited.
+class CtxWithSetup:  # safe: __aenter__ checkpoints, __aexit__ can be fast
+    async def __aenter__(self):
+        await foo()
+
+    async def __aexit__(self, exc_type, exc, tb):
+        print("fast exit")
+
+
+class CtxWithTeardown:  # safe: __aexit__ checkpoints, __aenter__ can be fast
+    async def __aenter__(self):
+        print("fast setup")
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await foo()
+
+
+class CtxWithBothCheckpoint:  # safe: both checkpoint
+    async def __aenter__(self):
+        await foo()
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await foo()
+
+
+# fmt: off
+class CtxNeitherCheckpoint:
+    async def __aenter__(self):  # error: 4, "exit", Stmt("function definition", line)
+        print("setup")
+
+    async def __aexit__(self, *a):  # error: 4, "exit", Stmt("function definition", line)
+        print("teardown")
+# fmt: on
+
+
+# Only one method defined: charitably assume the other is inherited with a checkpoint.
+class CtxOnlyAenter:  # safe: __aexit__ assumed inherited with checkpoint
+    async def __aenter__(self):
+        print("setup")
+
+
+class CtxOnlyAexit:  # safe: __aenter__ assumed inherited with checkpoint
+    async def __aexit__(self, *a):
+        print("teardown")
+
+
+class CtxOnlyAenterWithCheckpoint:  # safe
+    async def __aenter__(self):
+        await foo()
+
+
+class CtxOnlyAexitWithCheckpoint:  # safe
+    async def __aexit__(self, *a):
+        await foo()
+
+
+# a nested function named `__aenter__` inside another function is not a method
+def not_a_class():
+    async def __aenter__(self):  # error: 4, "exit", Stmt("function definition", line)
+        print("setup")
+
+
+# class nested inside a function still gets the exemption
+def factory():
+    class NestedCtx:  # safe
+        async def __aenter__(self):
+            print("setup")
+
+
+# nested class; outer class has nothing relevant
+class Outer:
+    class Inner:  # safe: charitable inheritance for __aexit__
+        async def __aenter__(self):
+            print("setup")