Handled padding only at boundaries of f-string astral-sh#4249

crates/ruff/resources/test/fixtures/flake8_comprehensions/C401.py

@@ -14,6 +14,8 @@ def f(x):
 print(f'Hello {set(a for a in "abc")} World')
 print(f"Hello {set(a for a in 'abc')} World")
 print(f"Hello {set(f(a) for a in 'abc')} World")
+print(f"{set(a for a in 'abc') - set(a for a in 'ab')}")
+print(f"{ set(a for a in 'abc') - set(a for a in 'ab') }")
 
 
 def set(*args, **kwargs):

crates/ruff/resources/test/fixtures/flake8_comprehensions/C402.py

@@ -9,6 +9,9 @@
 print(f'Hello {dict((x, x) for x in "abc")} World')
 print(f'Hello {dict((x,x) for x in "abc")} World')
 
+f'{dict((x, x) for x in range(3)) | dict((x, x) for x in range(3))}'
+f'{ dict((x, x) for x in range(3)) | dict((x, x) for x in range(3)) }'
+
 def f(x):
     return x
 

crates/ruff/resources/test/fixtures/flake8_comprehensions/C403.py

@@ -10,3 +10,6 @@ def f(x):
     return x
 
 s = f"{set([f(x) for x in 'ab'])}"
+
+s = f"{ set([x for x in 'ab']) | set([x for x in 'ab']) }"
+s = f"{set([x for x in 'ab']) | set([x for x in 'ab'])}"

crates/ruff/resources/test/fixtures/flake8_comprehensions/C404.py

@@ -8,3 +8,6 @@ def f(x):
 f"{dict([(s,s) for s in 'ab'])}"
 f"{dict([(s, s) for s in 'ab'])}"
 f"{dict([(s,f(s)) for s in 'ab'])}"
+
+f'{dict([(s,s) for s in "ab"]) | dict([(s,s) for s in "ab"])}'
+f'{ dict([(s,s) for s in "ab"]) | dict([(s,s) for s in "ab"]) }'

crates/ruff/resources/test/fixtures/flake8_comprehensions/C405.py

@@ -19,3 +19,6 @@
 f"{set([1,2,3])}"
 f"{set(['a', 'b'])}"
 f'{set(["a", "b"])}'
+
+f"{set(['a', 'b']) - set(['a'])}"
+f"{ set(['a', 'b']) - set(['a']) }"

crates/ruff/resources/test/fixtures/flake8_comprehensions/C408.py

@@ -14,3 +14,6 @@ def list():
 f"{dict(x='y')}"
 f'{dict(x="y")}'
 f"{dict()}"
+
+f"{dict(x='y') | dict(y='z')}"
+f"{ dict(x='y') | dict(y='z') }"

crates/ruff/src/checkers/ast/mod.rs

@@ -2892,22 +2892,12 @@ where
                 }
                 if self.settings.rules.enabled(Rule::UnnecessaryGeneratorSet) {
                     flake8_comprehensions::rules::unnecessary_generator_set(
-                        self,
-                        expr,
-                        self.ctx.expr_parent(),
-                        func,
-                        args,
-                        keywords,
+                        self, expr, func, args, keywords,
                     );
                 }
                 if self.settings.rules.enabled(Rule::UnnecessaryGeneratorDict) {
                     flake8_comprehensions::rules::unnecessary_generator_dict(
-                        self,
-                        expr,
-                        self.ctx.expr_parent(),
-                        func,
-                        args,
-                        keywords,
+                        self, expr, func, args, keywords,
                     );
                 }
                 if self

crates/ruff/src/rules/flake8_comprehensions/fixes.rs

@@ -7,6 +7,7 @@ use libcst_native::{
     ParenthesizedWhitespace, RightCurlyBrace, RightParen, RightSquareBracket, Set, SetComp,
     SimpleString, SimpleWhitespace, TrailingWhitespace, Tuple,
 };
+use ruff_text_size::{TextRange, TextSize};
 use rustpython_parser::ast::Ranged;
 
 use ruff_diagnostics::{Edit, Fix};
@@ -75,11 +76,12 @@ pub(crate) fn fix_unnecessary_generator_list(
 
 /// (C401) Convert `set(x for x in y)` to `{x for x in y}`.
 pub(crate) fn fix_unnecessary_generator_set(
-    locator: &Locator,
-    stylist: &Stylist,
+    checker: &Checker,
     expr: &rustpython_parser::ast::Expr,
-    parent: Option<&rustpython_parser::ast::Expr>,
 ) -> Result<Edit> {
+    let locator = checker.locator;
+    let stylist = checker.stylist;
+
     // Expr(Call(GeneratorExp)))) -> Expr(SetComp)))
     let module_text = locator.slice(expr.range());
     let mut tree = match_module(module_text)?;
@@ -113,25 +115,21 @@ pub(crate) fn fix_unnecessary_generator_set(
     };
     tree.codegen(&mut state);
 
-    let mut content = state.to_string();
-
-    // If the expression is embedded in an f-string, surround it with spaces to avoid
-    // syntax errors.
-    if let Some(rustpython_parser::ast::Expr::FormattedValue(_)) = parent {
-        content = format!(" {content} ");
-    }
-
-    Ok(Edit::range_replacement(content, expr.range()))
+    Ok(Edit::range_replacement(
+        wrap_code_in_spaces(&state, checker, expr.range()),
+        expr.range(),
+    ))
 }
 
 /// (C402) Convert `dict((x, x) for x in range(3))` to `{x: x for x in
 /// range(3)}`.
 pub(crate) fn fix_unnecessary_generator_dict(
-    locator: &Locator,
-    stylist: &Stylist,
+    checker: &Checker,
     expr: &rustpython_parser::ast::Expr,
-    parent: Option<&rustpython_parser::ast::Expr>,
 ) -> Result<Edit> {
+    let locator = checker.locator;
+    let stylist = checker.stylist;
+
     let module_text = locator.slice(expr.range());
     let mut tree = match_module(module_text)?;
     let mut body = match_expr(&mut tree)?;
@@ -181,15 +179,10 @@ pub(crate) fn fix_unnecessary_generator_dict(
     };
     tree.codegen(&mut state);
 
-    let mut content = state.to_string();
-
-    // If the expression is embedded in an f-string, surround it with spaces to avoid
-    // syntax errors.
-    if let Some(rustpython_parser::ast::Expr::FormattedValue(_)) = parent {
-        content = format!(" {content} ");
-    }
-
-    Ok(Edit::range_replacement(content, expr.range()))
+    Ok(Edit::range_replacement(
+        wrap_code_in_spaces(&state, checker, expr.range()),
+        expr.range(),
+    ))
 }
 
 /// (C403) Convert `set([x for x in y])` to `{x for x in y}`.
@@ -232,7 +225,7 @@ pub(crate) fn fix_unnecessary_list_comprehension_set(
     tree.codegen(&mut state);
 
     Ok(Edit::range_replacement(
-        wrap_code_in_spaces(&state, checker.ctx.in_f_string()),
+        wrap_code_in_spaces(&state, checker, expr.range()),
         expr.range(),
     ))
 }
@@ -288,7 +281,7 @@ pub(crate) fn fix_unnecessary_list_comprehension_dict(
     tree.codegen(&mut state);
 
     Ok(Edit::range_replacement(
-        wrap_code_in_spaces(&state, checker.ctx.in_f_string()),
+        wrap_code_in_spaces(&state, checker, expr.range()),
         expr.range(),
     ))
 }
@@ -383,17 +376,19 @@ pub(crate) fn fix_unnecessary_literal_set(
     tree.codegen(&mut state);
 
     Ok(Edit::range_replacement(
-        wrap_code_in_spaces(&state, checker.ctx.in_f_string()),
+        wrap_code_in_spaces(&state, checker, expr.range()),
         expr.range(),
     ))
 }
 
 /// (C406) Convert `dict([(1, 2)])` to `{1: 2}`.
 pub(crate) fn fix_unnecessary_literal_dict(
-    locator: &Locator,
-    stylist: &Stylist,
+    checker: &Checker,
     expr: &rustpython_parser::ast::Expr,
 ) -> Result<Edit> {
+    let locator = checker.locator;
+    let stylist = checker.stylist;
+
     // Expr(Call(List|Tuple)))) -> Expr(Dict)))
     let module_text = locator.slice(expr.range());
     let mut tree = match_module(module_text)?;
@@ -454,7 +449,10 @@ pub(crate) fn fix_unnecessary_literal_dict(
     };
     tree.codegen(&mut state);
 
-    Ok(Edit::range_replacement(state.to_string(), expr.range()))
+    Ok(Edit::range_replacement(
+        wrap_code_in_spaces(&state, checker, expr.range()),
+        expr.range(),
+    ))
 }
 
 /// (C408)
@@ -575,19 +573,51 @@ pub(crate) fn fix_unnecessary_collection_call(
     };
     tree.codegen(&mut state);
 
-    Ok(Edit::range_replacement(
-        wrap_code_in_spaces(&state, wrap_in_spaces),
-        expr.range(),
-    ))
+    let code = if wrap_in_spaces {
+        wrap_code_in_spaces(&state, checker, expr.range())
+    } else {
+        state.to_string()
+    };
+
+    Ok(Edit::range_replacement(code, expr.range()))
 }
 
 /// Adds spaces around the code generated from `state` if `wrap_in_spaces` is true.
-fn wrap_code_in_spaces(state: &CodegenState, wrap_in_spaces: bool) -> String {
-    if wrap_in_spaces {
-        format!(" {state} ")
-    } else {
-        state.to_string()
+fn wrap_code_in_spaces(state: &CodegenState, checker: &Checker, expr_range: TextRange) -> String {
+    if checker.ctx.in_f_string() {
+        if let Some(f_string_range) = checker.indexer.f_string_range(expr_range.start()) {
+            let f_string = checker.locator.slice(f_string_range);
+
+            // find the braces that delimit the f-string
+            let f_string_start = f_string_range.start();
+            let f_string_sbrace =
+                f_string_start + TextSize::try_from(f_string.find('{').unwrap()).unwrap();
+            let f_string_ebrace =
+                f_string_start + TextSize::try_from(f_string.rfind('}').unwrap()).unwrap();
+
+            let mut buf = Vec::with_capacity(3);
+
+            // check if left padding is required
+            // this is true when between the opening brace of the f-string and the start of the
+            // current expression there are no characters
+            let one = TextSize::from(1u32);
+            if f_string_sbrace + one == expr_range.start() {
+                buf.push(" ".to_string());
+            }
+
+            buf.push(state.to_string());
+
+            // check if right padding is required
+            // this is true when between the end of the current expression and the closing
+            // brace of the f-string there are no characters
+            if expr_range.end() == f_string_ebrace {
+                buf.push(" ".to_string());
+            }
+
+            return buf.join("");
+        }
     }
+    state.to_string()
 }
 
 /// (C409) Convert `tuple([1, 2])` to `tuple(1, 2)`

crates/ruff/src/rules/flake8_comprehensions/rules/unnecessary_generator_dict.rs

@@ -45,7 +45,6 @@ impl AlwaysAutofixableViolation for UnnecessaryGeneratorDict {
 pub(crate) fn unnecessary_generator_dict(
     checker: &mut Checker,
     expr: &Expr,
-    parent: Option<&Expr>,
     func: &Expr,
     args: &[Expr],
     keywords: &[Keyword],
@@ -60,12 +59,7 @@ pub(crate) fn unnecessary_generator_dict(
                 if checker.patch(diagnostic.kind.rule()) {
                     #[allow(deprecated)]
                     diagnostic.try_set_fix_from_edit(|| {
-                        fixes::fix_unnecessary_generator_dict(
-                            checker.locator,
-                            checker.stylist,
-                            expr,
-                            parent,
-                        )
+                        fixes::fix_unnecessary_generator_dict(checker, expr)
                     });
                 }
                 checker.diagnostics.push(diagnostic);

crates/ruff/src/rules/flake8_comprehensions/rules/unnecessary_generator_set.rs

@@ -45,7 +45,6 @@ impl AlwaysAutofixableViolation for UnnecessaryGeneratorSet {
 pub(crate) fn unnecessary_generator_set(
     checker: &mut Checker,
     expr: &Expr,
-    parent: Option<&Expr>,
     func: &Expr,
     args: &[Expr],
     keywords: &[Keyword],
@@ -60,9 +59,8 @@ pub(crate) fn unnecessary_generator_set(
         let mut diagnostic = Diagnostic::new(UnnecessaryGeneratorSet, expr.range());
         if checker.patch(diagnostic.kind.rule()) {
             #[allow(deprecated)]
-            diagnostic.try_set_fix_from_edit(|| {
-                fixes::fix_unnecessary_generator_set(checker.locator, checker.stylist, expr, parent)
-            });
+            diagnostic
+                .try_set_fix_from_edit(|| fixes::fix_unnecessary_generator_set(checker, expr));
         }
         checker.diagnostics.push(diagnostic);
     }

crates/ruff/src/rules/flake8_comprehensions/rules/unnecessary_literal_dict.rs

@@ -79,9 +79,7 @@ pub(crate) fn unnecessary_literal_dict(
     );
     if checker.patch(diagnostic.kind.rule()) {
         #[allow(deprecated)]
-        diagnostic.try_set_fix_from_edit(|| {
-            fixes::fix_unnecessary_literal_dict(checker.locator, checker.stylist, expr)
-        });
+        diagnostic.try_set_fix_from_edit(|| fixes::fix_unnecessary_literal_dict(checker, expr));
     }
     checker.diagnostics.push(diagnostic);
 }