From e4c2859c0f8908c70618603b8858fdd781e32b21 Mon Sep 17 00:00:00 2001
From: Dylan <53534755+dylwil3@users.noreply.github.com>
Date: Wed, 14 Aug 2024 20:02:57 -0500
Subject: [PATCH] [`flake8-async`] Do not lint yield in context manager
 `cancel-scope-no-checkpoint` (`ASYNC100`) (#12896)

For compatibility with upstream, treat `yield` as a checkpoint inside
cancel scopes.

Closes #12873.
---
 .../test/fixtures/flake8_async/ASYNC100.py    | 23 +++++++++++++++++++
 .../rules/cancel_scope_no_checkpoint.rs       | 15 ++++++++++--
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_async/ASYNC100.py b/crates/ruff_linter/resources/test/fixtures/flake8_async/ASYNC100.py
index 0ccdf30a0468d..8434073d22dac 100644
--- a/crates/ruff_linter/resources/test/fixtures/flake8_async/ASYNC100.py
+++ b/crates/ruff_linter/resources/test/fixtures/flake8_async/ASYNC100.py
@@ -89,3 +89,26 @@ async def func():
 async def func():
     async with asyncio.timeout(delay=0.2), asyncio.timeout(delay=0.2):
         ...
+
+
+# Don't trigger for blocks with a yield statement
+async def foo():
+    with trio.fail_after(1):
+        yield
+
+
+async def foo(): # even if only one branch contains a yield, we skip the lint
+    with trio.fail_after(1):
+        if something:
+            ...
+        else:
+            yield
+
+
+# https://github.com/astral-sh/ruff/issues/12873
+@asynccontextmanager
+async def good_code():
+    with anyio.fail_after(10):
+        # There's no await keyword here, but we presume that there
+        # will be in the caller we yield to, so this is safe.
+        yield
diff --git a/crates/ruff_linter/src/rules/flake8_async/rules/cancel_scope_no_checkpoint.rs b/crates/ruff_linter/src/rules/flake8_async/rules/cancel_scope_no_checkpoint.rs
index 26a5297ce3911..6b0b55b014654 100644
--- a/crates/ruff_linter/src/rules/flake8_async/rules/cancel_scope_no_checkpoint.rs
+++ b/crates/ruff_linter/src/rules/flake8_async/rules/cancel_scope_no_checkpoint.rs
@@ -1,8 +1,8 @@
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::helpers::AwaitVisitor;
+use ruff_python_ast::helpers::{any_over_body, AwaitVisitor};
 use ruff_python_ast::visitor::Visitor;
-use ruff_python_ast::{StmtWith, WithItem};
+use ruff_python_ast::{Expr, StmtWith, WithItem};
 
 use crate::checkers::ast::Checker;
 use crate::rules::flake8_async::helpers::MethodName;
@@ -10,6 +10,9 @@ use crate::rules::flake8_async::helpers::MethodName;
 /// ## What it does
 /// Checks for timeout context managers which do not contain a checkpoint.
 ///
+/// For the purposes of this check, `yield` is considered a checkpoint,
+/// since checkpoints may occur in the caller to which we yield.
+///
 /// ## Why is this bad?
 /// Some asynchronous context managers, such as `asyncio.timeout` and
 /// `trio.move_on_after`, have no effect unless they contain a checkpoint.
@@ -80,6 +83,14 @@ pub(crate) fn cancel_scope_no_checkpoint(
         return;
     }
 
+    // Treat yields as checkpoints, since checkpoints can happen
+    // in the caller yielded to.
+    // See: https://flake8-async.readthedocs.io/en/latest/rules.html#async100
+    // See: https://github.com/astral-sh/ruff/issues/12873
+    if any_over_body(&with_stmt.body, &Expr::is_yield_expr) {
+        return;
+    }
+
     // If the body contains an `await` statement, the context manager is used correctly.
     let mut visitor = AwaitVisitor::default();
     visitor.visit_body(&with_stmt.body);