Support parenthesized expressions in UP028 (#7114)

astral-sh · Sep 3, 2023 · b57ddd5 · b57ddd5
1 parent af189db
commit b57ddd5
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 2 deletions.
diff --git a/crates/ruff/resources/test/fixtures/pyupgrade/UP028_0.py b/crates/ruff/resources/test/fixtures/pyupgrade/UP028_0.py
@@ -72,3 +72,12 @@ def f():
     for x, y in z():
         yield x, y
     x = 1
+
+
+# Regression test for: https://github.com/astral-sh/ruff/issues/7103
+def _serve_method(fn):
+    for h in (
+        TaggedText.from_file(args.input)
+            .markup(highlight=args.region)
+    ):
+        yield h
diff --git a/crates/ruff/resources/test/fixtures/pyupgrade/UP028_1.py b/crates/ruff/resources/test/fixtures/pyupgrade/UP028_1.py
@@ -1,4 +1,4 @@
-# These should NOT change
+# OK
 def f():
     for x in z:
         yield

diff --git a/crates/ruff/src/rules/pyupgrade/rules/yield_in_for_loop.rs b/crates/ruff/src/rules/pyupgrade/rules/yield_in_for_loop.rs
@@ -1,5 +1,6 @@
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::parenthesize::parenthesized_range;
 use ruff_python_ast::{self as ast, Expr, Stmt};
 use ruff_text_size::Ranged;
 
@@ -103,7 +104,14 @@ pub(crate) fn yield_in_for_loop(checker: &mut Checker, stmt_for: &ast::StmtFor)
 
     let mut diagnostic = Diagnostic::new(YieldInForLoop, stmt_for.range());
     if checker.patch(diagnostic.kind.rule()) {
-        let contents = checker.locator().slice(iter.as_ref());
+        let contents = checker.locator().slice(
+            parenthesized_range(
+                iter.as_ref().into(),
+                stmt_for.into(),
+                checker.locator().contents(),
+            )
+            .unwrap_or(iter.range()),
+        );
         let contents = format!("yield from {contents}");
         diagnostic.set_fix(Fix::suggested(Edit::range_replacement(
             contents,

diff --git a/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP028_0.py.snap b/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP028_0.py.snap
@@ -270,5 +270,33 @@ UP028_0.py:72:5: UP028 [*] Replace `yield` over `for` loop with `yield from`
 73    |-        yield x, y
    72 |+    yield from z()
 74 73 |     x = 1
+75 74 | 
+76 75 | 
+
+UP028_0.py:79:5: UP028 [*] Replace `yield` over `for` loop with `yield from`
+   |
+77 |   # Regression test for: https://github.com/astral-sh/ruff/issues/7103
+78 |   def _serve_method(fn):
+79 |       for h in (
+   |  _____^
+80 | |         TaggedText.from_file(args.input)
+81 | |             .markup(highlight=args.region)
+82 | |     ):
+83 | |         yield h
+   | |_______________^ UP028
+   |
+   = help: Replace with `yield from`
+
+ℹ Suggested fix
+76 76 | 
+77 77 | # Regression test for: https://github.com/astral-sh/ruff/issues/7103
+78 78 | def _serve_method(fn):
+79    |-    for h in (
+   79 |+    yield from (
+80 80 |         TaggedText.from_file(args.input)
+81 81 |             .markup(highlight=args.region)
+82    |-    ):
+83    |-        yield h
+   82 |+    )