Tighten async CM exemption rules

Two refinements in response to review feedback:

- If an `__aenter__`/`__aexit__` method contains any checkpoint-like
  construct (`await`, `async with`, or `async for`), it must always
  checkpoint.  We no longer exempt such methods even when the partner
  provides a checkpoint -- conditional checkpoints are still flagged.

- Only charitably assume a missing partner is inherited (with a
  checkpoint) when the class actually inherits from something.  Classes
  with no base classes are treated as flat, and methods that don't
  checkpoint are flagged.  `metaclass=` and other keyword arguments do
  not count as inheriting, since they live in `ClassDef.keywords` rather
  than `ClassDef.bases`.

https://claude.ai/code/session_014jAydKywq31Ew4fVYGJdiG

Author: claude

flake8_async/visitors/visitor91x.py

@@ -467,11 +467,15 @@ def __init__(self, *args: Any, **kwargs: Any):
 
         # 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).
+        # that method contains a checkpoint-like construct, 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 inherited from a base class.
         self.async_cm_class: dict[str, bool] | None = None
+        # Whether the enclosing class has an explicit base class (other than
+        # implicit `object`). We only assume a missing partner is inherited if
+        # the class actually inherits from something.
+        self.async_cm_class_has_bases = False
         # Set on entry to an exempt `__aenter__`/`__aexit__` so that
         # `error_91x` skips emitting ASYNC910/911.
         self.exempt_async_cm_method = False
@@ -548,17 +552,24 @@ def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
     # 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")
+        self.save_state(node, "async_cm_class", "async_cm_class_has_bases")
         defined: dict[str, bool] = {}
+        checkpointy = (
+            m.Await()
+            | m.With(asynchronous=m.Asynchronous())
+            | m.For(asynchronous=m.Asynchronous())
+        )
         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()))
+                    defined[stmt.name.value] = bool(m.findall(stmt, checkpointy))
         self.async_cm_class = defined
+        # Keyword args like `metaclass=` are in `node.keywords`, not `bases`.
+        self.async_cm_class_has_bases = bool(node.bases)
 
     def leave_ClassDef(
         self, original_node: cst.ClassDef, updated_node: cst.ClassDef
@@ -574,11 +585,16 @@ def _is_exempt_async_cm_method(self, node: cst.FunctionDef) -> bool:
             return False
         if name not in self.async_cm_class:
             return False
+        # A method that contains any checkpoint must always checkpoint: we
+        # still check it normally so conditional checkpoints are flagged.
+        if self.async_cm_class[name]:
+            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.
+            # Partner is not defined on this class; only assume it is inherited
+            # (and contains a checkpoint) if the class inherits from something.
+            return self.async_cm_class_has_bases
+        # Partner defined; exempt iff it contains a checkpoint.
         return self.async_cm_class[partner]
 
     def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
@@ -609,6 +625,7 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
             "except_depth",
             "add_checkpoint_at_function_start",
             "async_cm_class",
+            "async_cm_class_has_bases",
             "exempt_async_cm_method",
             copy=True,
         )
@@ -622,6 +639,7 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
         self.add_checkpoint_at_function_start = False
         # Class-level context does not apply to nested scopes.
         self.async_cm_class = None
+        self.async_cm_class_has_bases = False
         self.exempt_async_cm_method = is_exempt_cm
 
         self.async_function = (

tests/autofix_files/async910.py

@@ -640,6 +640,10 @@ async def bar(): ...
 
 # Issue #441: async context manager methods may legitimately skip checkpointing
 # if the partner method provides the checkpoint, or if the partner is inherited.
+class ACM:  # a dummy base to opt into the charitable-inheritance assumption
+    pass
+
+
 class CtxWithSetup:  # safe: __aenter__ checkpoints, __aexit__ can be fast
     async def __aenter__(self):
         await foo()
@@ -676,17 +680,43 @@ async def __aexit__(self, *a):  # error: 4, "exit", Stmt("function definition",
 # fmt: on
 
 
-# Only one method defined: charitably assume the other is inherited with a checkpoint.
-class CtxOnlyAenter:  # safe: __aexit__ assumed inherited with checkpoint
+# A method that contains any checkpoint is still required to always checkpoint.
+class CtxAenterConditionalAexitFast(ACM):
+    async def __aenter__(self):  # error: 4, "exit", Stmt("function definition", line)
+        if _:
+            await foo()
+        await trio.lowlevel.checkpoint()
+
+    async def __aexit__(self, *a):
+        print("fast exit")
+
+
+# Only one method defined: charitably assume the other is inherited with a
+# checkpoint -- but only when the class inherits from something.
+class CtxOnlyAenterInherited(ACM):  # safe: __aexit__ assumed inherited
     async def __aenter__(self):
         print("setup")
 
 
-class CtxOnlyAexit:  # safe: __aenter__ assumed inherited with checkpoint
+class CtxOnlyAexitInherited(ACM):  # safe: __aenter__ assumed inherited
     async def __aexit__(self, *a):
         print("teardown")
 
 
+# fmt: off
+class CtxOnlyAenter:  # no base class -> don't assume inheritance
+    async def __aenter__(self):  # error: 4, "exit", Stmt("function definition", line)
+        print("setup")
+        await trio.lowlevel.checkpoint()
+
+
+class CtxOnlyAexit:  # no base class -> don't assume inheritance
+    async def __aexit__(self, *a):  # error: 4, "exit", Stmt("function definition", line)
+        print("teardown")
+        await trio.lowlevel.checkpoint()
+# fmt: on
+
+
 class CtxOnlyAenterWithCheckpoint:  # safe
     async def __aenter__(self):
         await foo()
@@ -697,22 +727,33 @@ async def __aexit__(self, *a):
         await foo()
 
 
+# keyword-only bases (like `metaclass=`) don't count as inheriting.
+class Meta(type):
+    pass
+
+
+class CtxMetaclassOnly(metaclass=Meta):
+    async def __aenter__(self):  # error: 4, "exit", Stmt("function definition", line)
+        print("setup")
+        await trio.lowlevel.checkpoint()
+
+
 # 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
+# class nested inside a function still gets the exemption when it inherits
 def factory():
-    class NestedCtx:  # safe
+    class NestedCtx(ACM):  # safe
         async def __aenter__(self):
             print("setup")
 
 
 # nested class; outer class has nothing relevant
 class Outer:
-    class Inner:  # safe: charitable inheritance for __aexit__
+    class Inner(ACM):  # safe: charitable inheritance for __aexit__
         async def __aenter__(self):
             print("setup")

tests/autofix_files/async910.py.diff

@@ -236,10 +236,39 @@
 
 
 @@ x,6 x,7 @@
+     async def __aenter__(self):  # error: 4, "exit", Stmt("function definition", line)
+         if _:
+             await foo()
++        await trio.lowlevel.checkpoint()
+
+     async def __aexit__(self, *a):
+         print("fast exit")
+@@ x,11 x,13 @@
+ class CtxOnlyAenter:  # no base class -> don't assume inheritance
+     async def __aenter__(self):  # error: 4, "exit", Stmt("function definition", line)
+         print("setup")
++        await trio.lowlevel.checkpoint()
+
+
+ class CtxOnlyAexit:  # no base class -> don't assume inheritance
+     async def __aexit__(self, *a):  # error: 4, "exit", Stmt("function definition", line)
+         print("teardown")
++        await trio.lowlevel.checkpoint()
+ # fmt: on
+
+
+@@ x,12 x,14 @@
+ class CtxMetaclassOnly(metaclass=Meta):
+     async def __aenter__(self):  # error: 4, "exit", Stmt("function definition", line)
+         print("setup")
++        await trio.lowlevel.checkpoint()
+
+
+ # 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
+ # class nested inside a function still gets the exemption when it inherits

tests/eval_files/async910.py

@@ -610,6 +610,10 @@ async def bar(): ...
 
 # Issue #441: async context manager methods may legitimately skip checkpointing
 # if the partner method provides the checkpoint, or if the partner is inherited.
+class ACM:  # a dummy base to opt into the charitable-inheritance assumption
+    pass
+
+
 class CtxWithSetup:  # safe: __aenter__ checkpoints, __aexit__ can be fast
     async def __aenter__(self):
         await foo()
@@ -644,17 +648,40 @@ async def __aexit__(self, *a):  # error: 4, "exit", Stmt("function definition",
 # fmt: on
 
 
-# Only one method defined: charitably assume the other is inherited with a checkpoint.
-class CtxOnlyAenter:  # safe: __aexit__ assumed inherited with checkpoint
+# A method that contains any checkpoint is still required to always checkpoint.
+class CtxAenterConditionalAexitFast(ACM):
+    async def __aenter__(self):  # error: 4, "exit", Stmt("function definition", line)
+        if _:
+            await foo()
+
+    async def __aexit__(self, *a):
+        print("fast exit")
+
+
+# Only one method defined: charitably assume the other is inherited with a
+# checkpoint -- but only when the class inherits from something.
+class CtxOnlyAenterInherited(ACM):  # safe: __aexit__ assumed inherited
     async def __aenter__(self):
         print("setup")
 
 
-class CtxOnlyAexit:  # safe: __aenter__ assumed inherited with checkpoint
+class CtxOnlyAexitInherited(ACM):  # safe: __aenter__ assumed inherited
     async def __aexit__(self, *a):
         print("teardown")
 
 
+# fmt: off
+class CtxOnlyAenter:  # no base class -> don't assume inheritance
+    async def __aenter__(self):  # error: 4, "exit", Stmt("function definition", line)
+        print("setup")
+
+
+class CtxOnlyAexit:  # no base class -> don't assume inheritance
+    async def __aexit__(self, *a):  # error: 4, "exit", Stmt("function definition", line)
+        print("teardown")
+# fmt: on
+
+
 class CtxOnlyAenterWithCheckpoint:  # safe
     async def __aenter__(self):
         await foo()
@@ -665,21 +692,31 @@ async def __aexit__(self, *a):
         await foo()
 
 
+# keyword-only bases (like `metaclass=`) don't count as inheriting.
+class Meta(type):
+    pass
+
+
+class CtxMetaclassOnly(metaclass=Meta):
+    async def __aenter__(self):  # error: 4, "exit", Stmt("function definition", line)
+        print("setup")
+
+
 # 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
+# class nested inside a function still gets the exemption when it inherits
 def factory():
-    class NestedCtx:  # safe
+    class NestedCtx(ACM):  # safe
         async def __aenter__(self):
             print("setup")
 
 
 # nested class; outer class has nothing relevant
 class Outer:
-    class Inner:  # safe: charitable inheritance for __aexit__
+    class Inner(ACM):  # safe: charitable inheritance for __aexit__
         async def __aenter__(self):
             print("setup")