add async124 yield-in-asynccm-not-in-try

Author: jakkdl

docs/rules.rst

@@ -94,6 +94,9 @@ _`ASYNC123`: bad-exception-group-flattening
     Dropping this information makes diagnosing errors much more difficult.
     We recommend ``raise SomeNewError(...) from group`` if possible; or consider using `copy.copy` to shallow-copy the exception before re-raising (for copyable types), or re-raising the error from outside the `except` block.
 
+_`ASYNC124`: yield-in-asynccm-not-in-try
+    `yield` in ``@asynccontextmanager`` should usually be in a ``try:`` with cleanup code in ``finally:`` so cleanup runs if the yielded code raises an exception.
+
 Blocking sync calls in async functions
 ======================================
 

flake8_async/visitors/visitor102.py

@@ -161,6 +161,9 @@ def visit_Try(self, node: ast.Try):
         self._critical_scope = Statement("try/finally", node.lineno, node.col_offset)
         self.visit_nodes(node.finalbody)
 
+        # don't revisit children
+        self.novisit = True
+
     def visit_ExceptHandler(self, node: ast.ExceptHandler):
         # if we're inside a critical scope, a nested except should never override that
         if self._critical_scope is not None and self._critical_scope.name != "except":

flake8_async/visitors/visitors.py

@@ -437,6 +437,48 @@ def visit_Call(self, node: ast.Call):
             self.error(node, f"{match[2]}.{match[1]}")
 
 
+@error_class
+class Visitor124(Flake8AsyncVisitor):
+    error_codes: Mapping[str, str] = {
+        "ASYNC124": (
+            "yield in @asynccontextmanager should usually be in a `try:`"
+            " with cleanup in `finally` to ensure cleanup is run."
+        )
+    }
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        self.in_asynccontextmanager: bool = False
+        self.in_try: bool = False
+
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
+        self.save_state(node, "in_asynccontextmanager")
+
+        # TODO: should it also error on @contextmanager? If so it feels like
+        # flake8-bugbear might be more appropriate.
+        if has_decorator(node, "asynccontextmanager"):
+            self.in_asynccontextmanager = True
+
+    # ast.TryStar added in py311, we run mypy on py39
+    def visit_Try(self, node: ast.Try | ast.TryStar):  # type: ignore[name-defined]
+        old_in_try = self.in_try
+        self.in_try = True
+
+        self.visit_nodes(node.body)
+        self.in_try = old_in_try
+
+        self.visit_nodes(node.handlers, node.orelse, node.finalbody)
+
+        # don't revisit children
+        self.novisit = True
+
+    visit_TryStar = visit_Try
+
+    def visit_Yield(self, node: ast.Yield):
+        if self.in_asynccontextmanager and not self.in_try:
+            self.error(node)
+
+
 @error_class_cst
 class Visitor300(Flake8AsyncVisitor_cst):
     error_codes: Mapping[str, str] = {

tests/eval_files/async124.py

@@ -0,0 +1,18 @@
+from contextlib import asynccontextmanager
+
+
+@asynccontextmanager
+async def foo():
+    try:
+        # TODO: should it error if there is no finally?
+        yield
+    except:
+        ...
+
+
+@asynccontextmanager
+async def foo2():
+    try:
+        ...
+    except:
+        yield  # error: 8

tests/eval_files/async124_py311.py

@@ -0,0 +1,17 @@
+from contextlib import asynccontextmanager
+
+
+@asynccontextmanager
+async def foo():
+    try:
+        yield
+    except* Exception:
+        ...
+
+
+@asynccontextmanager
+async def bar():
+    try:
+        ...
+    except* Exception:
+        yield  # error: 8