From f7dcad399b998c7a40b54d53db9bb6b0748952cb Mon Sep 17 00:00:00 2001
From: jakkdl <h6+github@pm.me>
Date: Wed, 24 Apr 2024 10:36:35 +0200
Subject: [PATCH] make async900 label point directly to the async900 rule. Save
 bools instead of nodes since we don't need to warn on them or refer to them.
 Warn on each yield, add test case for that.

---
 docs/rules.rst                    |  3 ++-
 flake8_async/visitors/visitors.py | 20 +++++++++-----------
 tests/eval_files/async119.py      |  9 +++++++++
 3 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/docs/rules.rst b/docs/rules.rst
index eabaf54..f1eea69 100644
--- a/docs/rules.rst
+++ b/docs/rules.rst
@@ -45,11 +45,12 @@ Note: 22X, 23X and 24X has not had asyncio-specific suggestions written.
 - **ASYNC250**: Builtin ``input()`` should not be called from async function. Wrap in ``[trio/anyio].to_thread.run_sync()`` or ``asyncio.loop.run_in_executor()``.
 - **ASYNC251**: ``time.sleep(...)`` should not be called from async function. Use ``[trio/anyio/asyncio].sleep(...)``.
 
-.. _async900:
 
 Optional rules disabled by default
 ==================================
 
+.. _async900:
+
 - **ASYNC900**: Async generator without ``@asynccontextmanager`` not allowed. You might want to enable this on a codebase since async generators are inherently unsafe and cleanup logic might not be performed. See https://github.com/python-trio/flake8-async/issues/211 and https://discuss.python.org/t/using-exceptiongroup-at-anthropic-experience-report/20888/6 for discussion.
 - **ASYNC910**: Exit or ``return`` from async function with no guaranteed checkpoint or exception since function definition. You might want to enable this on a codebase to make it easier to reason about checkpoints, and make the logic of ASYNC911 correct.
 - **ASYNC911**: Exit, ``yield`` or ``return`` from async iterable with no guaranteed checkpoint since possible function entry (yield or function definition)
diff --git a/flake8_async/visitors/visitors.py b/flake8_async/visitors/visitors.py
index de91e2f..b947ddb 100644
--- a/flake8_async/visitors/visitors.py
+++ b/flake8_async/visitors/visitors.py
@@ -290,35 +290,33 @@ class Visitor119(Flake8AsyncVisitor):
 
     def __init__(self, *args: Any, **kwargs: Any):
         super().__init__(*args, **kwargs)
-        self.unsafe_function: ast.AsyncFunctionDef | None = None
-        self.contextmanager: ast.With | ast.AsyncWith | None = None
+        self.unsafe_function: bool = False
+        self.contextmanager: bool = False
 
     def visit_AsyncFunctionDef(
         self, node: ast.AsyncFunctionDef | ast.FunctionDef | ast.Lambda
     ):
         self.save_state(node, "unsafe_function", "contextmanager")
-        self.contextmanager = None
+        self.contextmanager = False
         if isinstance(node, ast.AsyncFunctionDef) and not has_decorator(
             node, "asynccontextmanager"
         ):
-            self.unsafe_function = node
+            self.unsafe_function = True
         else:
-            self.unsafe_function = None
+            self.unsafe_function = False
 
     def visit_With(self, node: ast.With | ast.AsyncWith):
         self.save_state(node, "contextmanager")
-        self.contextmanager = node
+        self.contextmanager = True
 
     def visit_Yield(self, node: ast.Yield):
-        if self.unsafe_function is not None and self.contextmanager is not None:
-            # Decision point: the error could point to the method, or context manager,
-            # or the yield.
+        if self.unsafe_function and self.contextmanager:
             self.error(node)
-            # only warn once per method (?)
-            self.unsafe_function = None
 
     visit_AsyncWith = visit_With
     visit_FunctionDef = visit_AsyncFunctionDef
+    # it's not possible to yield or open context managers in lambda's, so this
+    # one isn't strictly needed afaik.
     visit_Lambda = visit_AsyncFunctionDef
 
 
diff --git a/tests/eval_files/async119.py b/tests/eval_files/async119.py
index 3f173c4..242db2c 100644
--- a/tests/eval_files/async119.py
+++ b/tests/eval_files/async119.py
@@ -13,6 +13,15 @@ async def async_with():
         yield  # error: 8
 
 
+async def warn_on_yeach_yield():
+    with open(""):
+        yield  # error: 8
+        yield  # error: 8
+    with open(""):
+        yield  # error: 8
+        yield  # error: 8
+
+
 async def yield_not_in_contextmanager():
     yield
     with open(""):