[`refurb`] Parenthesize lambda and ternary expressions in iter (`FURB122`, `FURB142`) (#18592)

Summary
--

Fixes #18590 by adding parentheses around lambdas and if expressions in
`for` loop iterators for FURB122 and FURB142. I also updated the docs on
the helper function to reflect the part actually being parenthesized and
the new checks.

The `lambda` case actually causes a `TypeError` at runtime, but I think
it's still worth handling to avoid causing a syntax error.

```pycon
>>> s = set()
... for x in (1,) if True else (2,):
...     s.add(-x)
... for x in lambda: 0:
...     s.discard(-x)
...
Traceback (most recent call last):
  File "<python-input-0>", line 4, in <module>
    for x in lambda: 0:
             ^^^^^^^^^
TypeError: 'function' object is not iterable
```

Test Plan
--

New test cases based on the bug report

---------

Co-authored-by: Dylan <dylwil3@gmail.com>

crates/ruff_linter/resources/test/fixtures/refurb/FURB122.py

@@ -174,3 +174,43 @@ def _():
         global global_foo
         for [a, b, (global_foo, c)] in d:
             f.write((a, b))
+
+
+# Test cases for lambda and ternary expressions - https://github.com/astral-sh/ruff/issues/18590
+
+def _():
+    with Path("file.txt").open("w", encoding="utf-8") as f:
+        for l in lambda: 0:
+            f.write(f"[{l}]")
+
+
+def _():
+    with Path("file.txt").open("w", encoding="utf-8") as f:
+        for l in (1,) if True else (2,):
+            f.write(f"[{l}]")
+
+
+# don't need to add parentheses when making a function argument
+def _():
+    with open("file", "w") as f:
+        for line in lambda: 0:
+            f.write(line)
+
+
+def _():
+    with open("file", "w") as f:
+        for line in (1,) if True else (2,):
+            f.write(line)
+
+
+# don't add extra parentheses if they're already parenthesized
+def _():
+    with open("file", "w") as f:
+        for line in (lambda: 0):
+            f.write(f"{line}")
+
+
+def _():
+    with open("file", "w") as f:
+        for line in ((1,) if True else (2,)):
+            f.write(f"{line}")

crates/ruff_linter/resources/test/fixtures/refurb/FURB142.py

@@ -74,3 +74,28 @@ async def f(y):
 def g():
     for x in (set(),):
         x.add(x)
+
+
+# Test cases for lambda and ternary expressions - https://github.com/astral-sh/ruff/issues/18590
+
+s = set()
+
+for x in lambda: 0:
+    s.discard(-x)
+
+for x in (1,) if True else (2,):
+    s.add(-x)
+
+# don't add extra parens
+for x in (lambda: 0):
+    s.discard(-x)
+
+for x in ((1,) if True else (2,)):
+    s.add(-x)
+
+# don't add parens directly in function call
+for x in lambda: 0:
+    s.discard(x)
+
+for x in (1,) if True else (2,):
+    s.add(x)

crates/ruff_linter/src/rules/refurb/rules/for_loop_set_mutations.rs

@@ -3,6 +3,7 @@ use ruff_python_ast::{Expr, Stmt, StmtFor};
 use ruff_python_semantic::analyze::typing;
 
 use crate::checkers::ast::Checker;
+use crate::rules::refurb::rules::helpers::IterLocation;
 use crate::{AlwaysFixableViolation, Applicability, Edit, Fix};
 
 use super::helpers::parenthesize_loop_iter_if_necessary;
@@ -106,7 +107,7 @@ pub(crate) fn for_loop_set_mutations(checker: &Checker, for_stmt: &StmtFor) {
             format!(
                 "{}.{batch_method_name}({})",
                 set.id,
-                parenthesize_loop_iter_if_necessary(for_stmt, checker),
+                parenthesize_loop_iter_if_necessary(for_stmt, checker, IterLocation::Call),
             )
         }
         (for_target, arg) => format!(
@@ -114,7 +115,7 @@ pub(crate) fn for_loop_set_mutations(checker: &Checker, for_stmt: &StmtFor) {
             set.id,
             locator.slice(arg),
             locator.slice(for_target),
-            parenthesize_loop_iter_if_necessary(for_stmt, checker),
+            parenthesize_loop_iter_if_necessary(for_stmt, checker, IterLocation::Comprehension),
         ),
     };
 

crates/ruff_linter/src/rules/refurb/rules/for_loop_writes.rs

@@ -5,6 +5,7 @@ use ruff_python_semantic::{Binding, ScopeId, SemanticModel, TypingOnlyBindingsSt
 use ruff_text_size::{Ranged, TextRange, TextSize};
 
 use crate::checkers::ast::Checker;
+use crate::rules::refurb::rules::helpers::IterLocation;
 use crate::{AlwaysFixableViolation, Applicability, Edit, Fix};
 
 use super::helpers::parenthesize_loop_iter_if_necessary;
@@ -182,7 +183,7 @@ fn for_loop_writes(
             format!(
                 "{}.writelines({})",
                 locator.slice(io_object_name),
-                parenthesize_loop_iter_if_necessary(for_stmt, checker),
+                parenthesize_loop_iter_if_necessary(for_stmt, checker, IterLocation::Call),
             )
         }
         (for_target, write_arg) => {
@@ -191,7 +192,7 @@ fn for_loop_writes(
                 locator.slice(io_object_name),
                 locator.slice(write_arg),
                 locator.slice(for_target),
-                parenthesize_loop_iter_if_necessary(for_stmt, checker),
+                parenthesize_loop_iter_if_necessary(for_stmt, checker, IterLocation::Comprehension),
             )
         }
     };

crates/ruff_linter/src/rules/refurb/rules/helpers.rs

@@ -4,16 +4,24 @@ use ruff_python_ast::{self as ast, parenthesize::parenthesized_range};
 
 use crate::checkers::ast::Checker;
 
-/// A helper function that extracts the `iter` from a [`ast::StmtFor`] node and,
+/// A helper function that extracts the `iter` from a [`ast::StmtFor`] node and
-/// if the `iter` is an unparenthesized tuple, adds parentheses:
+/// adds parentheses if needed.
 ///
-/// - `for x in z: ...`       ->  `"x"`
+/// These cases are okay and will not be modified:
-/// - `for (x, y) in z: ...`  ->  `"(x, y)"`
+///
-/// - `for [x, y] in z: ...`  ->  `"[x, y]"`
+/// - `for x in z: ...`       ->  `"z"`
-/// - `for x, y in z: ...`    ->  `"(x, y)"`      # <-- Parentheses added only for this example
+/// - `for x in (y, z): ...`  ->  `"(y, z)"`
+/// - `for x in [y, z]: ...`  ->  `"[y, z]"`
+///
+/// While these cases require parentheses:
+///
+/// - `for x in y, z: ...`                   ->  `"(y, z)"`
+/// - `for x in lambda: 0: ...`              ->  `"(lambda: 0)"`
+/// - `for x in (1,) if True else (2,): ...` ->  `"((1,) if True else (2,))"`
 pub(super) fn parenthesize_loop_iter_if_necessary<'a>(
     for_stmt: &'a ast::StmtFor,
     checker: &'a Checker,
+    location: IterLocation,
 ) -> Cow<'a, str> {
     let locator = checker.locator();
     let iter = for_stmt.iter.as_ref();
@@ -35,6 +43,17 @@ pub(super) fn parenthesize_loop_iter_if_necessary<'a>(
         ast::Expr::Tuple(tuple) if !tuple.parenthesized => {
             Cow::Owned(format!("({iter_in_source})"))
         }
+        ast::Expr::Lambda(_) | ast::Expr::If(_)
+            if matches!(location, IterLocation::Comprehension) =>
+        {
+            Cow::Owned(format!("({iter_in_source})"))
+        }
         _ => Cow::Borrowed(iter_in_source),
     }
 }
+
+#[derive(Copy, Clone)]
+pub(super) enum IterLocation {
+    Call,
+    Comprehension,
+}

crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB122_FURB122.py.snap

@@ -307,3 +307,126 @@ FURB122.py:93:9: FURB122 [*] Use of `f.write` in a for loop
 98 97 | 
 99 98 | 
 100 99 | # OK
+
+FURB122.py:183:9: FURB122 [*] Use of `f.write` in a for loop
+    |
+181 |   def _():
+182 |       with Path("file.txt").open("w", encoding="utf-8") as f:
+183 | /         for l in lambda: 0:
+184 | |             f.write(f"[{l}]")
+    | |_____________________________^ FURB122
+    |
+    = help: Replace with `f.writelines`
+
+ℹ Safe fix
+180 180 | 
+181 181 | def _():
+182 182 |     with Path("file.txt").open("w", encoding="utf-8") as f:
+183     |-        for l in lambda: 0:
+184     |-            f.write(f"[{l}]")
+    183 |+        f.writelines(f"[{l}]" for l in (lambda: 0))
+185 184 | 
+186 185 | 
+187 186 | def _():
+
+FURB122.py:189:9: FURB122 [*] Use of `f.write` in a for loop
+    |
+187 |   def _():
+188 |       with Path("file.txt").open("w", encoding="utf-8") as f:
+189 | /         for l in (1,) if True else (2,):
+190 | |             f.write(f"[{l}]")
+    | |_____________________________^ FURB122
+    |
+    = help: Replace with `f.writelines`
+
+ℹ Safe fix
+186 186 | 
+187 187 | def _():
+188 188 |     with Path("file.txt").open("w", encoding="utf-8") as f:
+189     |-        for l in (1,) if True else (2,):
+190     |-            f.write(f"[{l}]")
+    189 |+        f.writelines(f"[{l}]" for l in ((1,) if True else (2,)))
+191 190 | 
+192 191 | 
+193 192 | # don't need to add parentheses when making a function argument
+
+FURB122.py:196:9: FURB122 [*] Use of `f.write` in a for loop
+    |
+194 |   def _():
+195 |       with open("file", "w") as f:
+196 | /         for line in lambda: 0:
+197 | |             f.write(line)
+    | |_________________________^ FURB122
+    |
+    = help: Replace with `f.writelines`
+
+ℹ Safe fix
+193 193 | # don't need to add parentheses when making a function argument
+194 194 | def _():
+195 195 |     with open("file", "w") as f:
+196     |-        for line in lambda: 0:
+197     |-            f.write(line)
+    196 |+        f.writelines(lambda: 0)
+198 197 | 
+199 198 | 
+200 199 | def _():
+
+FURB122.py:202:9: FURB122 [*] Use of `f.write` in a for loop
+    |
+200 |   def _():
+201 |       with open("file", "w") as f:
+202 | /         for line in (1,) if True else (2,):
+203 | |             f.write(line)
+    | |_________________________^ FURB122
+    |
+    = help: Replace with `f.writelines`
+
+ℹ Safe fix
+199 199 | 
+200 200 | def _():
+201 201 |     with open("file", "w") as f:
+202     |-        for line in (1,) if True else (2,):
+203     |-            f.write(line)
+    202 |+        f.writelines((1,) if True else (2,))
+204 203 | 
+205 204 | 
+206 205 | # don't add extra parentheses if they're already parenthesized
+
+FURB122.py:209:9: FURB122 [*] Use of `f.write` in a for loop
+    |
+207 |   def _():
+208 |       with open("file", "w") as f:
+209 | /         for line in (lambda: 0):
+210 | |             f.write(f"{line}")
+    | |______________________________^ FURB122
+    |
+    = help: Replace with `f.writelines`
+
+ℹ Safe fix
+206 206 | # don't add extra parentheses if they're already parenthesized
+207 207 | def _():
+208 208 |     with open("file", "w") as f:
+209     |-        for line in (lambda: 0):
+210     |-            f.write(f"{line}")
+    209 |+        f.writelines(f"{line}" for line in (lambda: 0))
+211 210 | 
+212 211 | 
+213 212 | def _():
+
+FURB122.py:215:9: FURB122 [*] Use of `f.write` in a for loop
+    |
+213 |   def _():
+214 |       with open("file", "w") as f:
+215 | /         for line in ((1,) if True else (2,)):
+216 | |             f.write(f"{line}")
+    | |______________________________^ FURB122
+    |
+    = help: Replace with `f.writelines`
+
+ℹ Safe fix
+212 212 | 
+213 213 | def _():
+214 214 |     with open("file", "w") as f:
+215     |-        for line in ((1,) if True else (2,)):
+216     |-            f.write(f"{line}")
+    215 |+        f.writelines(f"{line}" for line in ((1,) if True else (2,)))

crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB142_FURB142.py.snap

@@ -280,3 +280,134 @@ FURB142.py:41:1: FURB142 [*] Use of `set.add()` in a for loop
 46 45 | 
 47 46 | 
 48 47 | # False negative
+
+FURB142.py:83:1: FURB142 [*] Use of `set.discard()` in a for loop
+   |
+81 |   s = set()
+82 |
+83 | / for x in lambda: 0:
+84 | |     s.discard(-x)
+   | |_________________^ FURB142
+85 |
+86 |   for x in (1,) if True else (2,):
+   |
+   = help: Replace with `.difference_update()`
+
+ℹ Safe fix
+80 80 | 
+81 81 | s = set()
+82 82 | 
+83    |-for x in lambda: 0:
+84    |-    s.discard(-x)
+   83 |+s.difference_update(-x for x in (lambda: 0))
+85 84 | 
+86 85 | for x in (1,) if True else (2,):
+87 86 |     s.add(-x)
+
+FURB142.py:86:1: FURB142 [*] Use of `set.add()` in a for loop
+   |
+84 |       s.discard(-x)
+85 |
+86 | / for x in (1,) if True else (2,):
+87 | |     s.add(-x)
+   | |_____________^ FURB142
+88 |
+89 |   # don't add extra parens
+   |
+   = help: Replace with `.update()`
+
+ℹ Safe fix
+83 83 | for x in lambda: 0:
+84 84 |     s.discard(-x)
+85 85 | 
+86    |-for x in (1,) if True else (2,):
+87    |-    s.add(-x)
+   86 |+s.update(-x for x in ((1,) if True else (2,)))
+88 87 | 
+89 88 | # don't add extra parens
+90 89 | for x in (lambda: 0):
+
+FURB142.py:90:1: FURB142 [*] Use of `set.discard()` in a for loop
+   |
+89 |   # don't add extra parens
+90 | / for x in (lambda: 0):
+91 | |     s.discard(-x)
+   | |_________________^ FURB142
+92 |
+93 |   for x in ((1,) if True else (2,)):
+   |
+   = help: Replace with `.difference_update()`
+
+ℹ Safe fix
+87 87 |     s.add(-x)
+88 88 | 
+89 89 | # don't add extra parens
+90    |-for x in (lambda: 0):
+91    |-    s.discard(-x)
+   90 |+s.difference_update(-x for x in (lambda: 0))
+92 91 | 
+93 92 | for x in ((1,) if True else (2,)):
+94 93 |     s.add(-x)
+
+FURB142.py:93:1: FURB142 [*] Use of `set.add()` in a for loop
+   |
+91 |       s.discard(-x)
+92 |
+93 | / for x in ((1,) if True else (2,)):
+94 | |     s.add(-x)
+   | |_____________^ FURB142
+95 |
+96 |   # don't add parens directly in function call
+   |
+   = help: Replace with `.update()`
+
+ℹ Safe fix
+90 90 | for x in (lambda: 0):
+91 91 |     s.discard(-x)
+92 92 | 
+93    |-for x in ((1,) if True else (2,)):
+94    |-    s.add(-x)
+   93 |+s.update(-x for x in ((1,) if True else (2,)))
+95 94 | 
+96 95 | # don't add parens directly in function call
+97 96 | for x in lambda: 0:
+
+FURB142.py:97:1: FURB142 [*] Use of `set.discard()` in a for loop
+    |
+ 96 |   # don't add parens directly in function call
+ 97 | / for x in lambda: 0:
+ 98 | |     s.discard(x)
+    | |________________^ FURB142
+ 99 |
+100 |   for x in (1,) if True else (2,):
+    |
+    = help: Replace with `.difference_update()`
+
+ℹ Safe fix
+94 94 |     s.add(-x)
+95 95 | 
+96 96 | # don't add parens directly in function call
+97    |-for x in lambda: 0:
+98    |-    s.discard(x)
+   97 |+s.difference_update(lambda: 0)
+99 98 | 
+100 99 | for x in (1,) if True else (2,):
+101 100 |     s.add(x)
+
+FURB142.py:100:1: FURB142 [*] Use of `set.add()` in a for loop
+    |
+ 98 |       s.discard(x)
+ 99 |
+100 | / for x in (1,) if True else (2,):
+101 | |     s.add(x)
+    | |____________^ FURB142
+    |
+    = help: Replace with `.update()`
+
+ℹ Safe fix
+97  97  | for x in lambda: 0:
+98  98  |     s.discard(x)
+99  99  | 
+100     |-for x in (1,) if True else (2,):
+101     |-    s.add(x)
+    100 |+s.update((1,) if True else (2,))