Handled dict and set inside f-string (astral-sh#4249)

This fixes some C40* rules by adding spaces around dict and
set literals inside f-string, and quotes correctly string literals
inside f-string, since dict kwargs must be converted into
string literals.

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

@@ -6,6 +6,16 @@
 _ = '{}'.format(set(a if a < 6 else 0  for a in range(3)))
 print(f'Hello {set(a for a in range(3))} World')
 
+
+def f(x):
+    return 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")
+
+
 def set(*args, **kwargs):
     return None
 

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

@@ -5,3 +5,11 @@
 dict(((x, x) for x in range(3)), z=3)
 y = f'{dict((x, x) for x in range(3))}'
 print(f'Hello {dict((x, x) for x in range(3))} World')
+print(f"Hello {dict((x, x) for x in 'abc')} World")
+print(f'Hello {dict((x, x) for x in "abc")} World')
+print(f'Hello {dict((x,x) for x in "abc")} World')
+
+def f(x):
+    return x
+
+print(f'Hello {dict((x,f(x)) for x in "abc")} World')

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

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

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

@@ -1,2 +1,10 @@
 dict([(i, i) for i in range(3)])
 dict([(i, i) for i in range(3)], z=4)
+
+def f(x):
+    return x
+
+f'{dict([(s,s) for s in "ab"])}'
+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'])}"

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

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

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

@@ -10,3 +10,7 @@ def list():
 
 
 a = list()
+
+f"{dict(x='y')}"
+f'{dict(x="y")}'
+f"{dict()}"

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

@@ -138,30 +138,35 @@ impl<'a> Checker<'a> {
 
     /// Create a [`Generator`] to generate source code based on the current AST state.
     pub(crate) fn generator(&self) -> Generator {
-        fn quote_style(context: &Context, locator: &Locator, indexer: &Indexer) -> Option<Quote> {
-            if !context.in_f_string() {
-                return None;
-            }
-
-            // Find the quote character used to start the containing f-string.
-            let expr = context.expr()?;
-            let string_range = indexer.f_string_range(expr.start())?;
-            let trailing_quote = trailing_quote(locator.slice(string_range))?;
-
-            // Invert the quote character, if it's a single quote.
-            match *trailing_quote {
-                "'" => Some(Quote::Double),
-                "\"" => Some(Quote::Single),
-                _ => None,
-            }
-        }
-
         Generator::new(
             self.stylist.indentation(),
-            quote_style(&self.ctx, self.locator, self.indexer).unwrap_or(self.stylist.quote()),
+            self.f_string_quote_style().unwrap_or(self.stylist.quote()),
             self.stylist.line_ending(),
         )
     }
+
+    /// Returns the appropriate quoting for f-string by reversing the one used outside of
+    /// the f-string.
+    ///
+    /// If the current expression in the context is not an f-string, returns ``None``.
+    pub(crate) fn f_string_quote_style(&self) -> Option<Quote> {
+        let context = &self.ctx;
+        if !context.in_f_string() {
+            return None;
+        }
+
+        // Find the quote character used to start the containing f-string.
+        let expr = context.expr()?;
+        let string_range = self.indexer.f_string_range(expr.start())?;
+        let trailing_quote = trailing_quote(self.locator.slice(string_range))?;
+
+        // Invert the quote character, if it's a single quote.
+        match *trailing_quote {
+            "'" => Some(Quote::Double),
+            "\"" => Some(Quote::Single),
+            _ => None,
+        }
+    }
 }
 
 impl<'a, 'b> Visitor<'b> for Checker<'a>

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

@@ -12,6 +12,7 @@ use rustpython_parser::ast::Ranged;
 use ruff_diagnostics::{Edit, Fix};
 use ruff_python_ast::source_code::{Locator, Stylist};
 
+use crate::checkers::ast::Checker;
 use crate::cst::matchers::{match_expr, match_module};
 
 fn match_call<'a, 'b>(expr: &'a mut Expr<'b>) -> Result<&'a mut Call<'b>> {
@@ -193,10 +194,11 @@ pub(crate) fn fix_unnecessary_generator_dict(
 
 /// (C403) Convert `set([x for x in y])` to `{x for x in y}`.
 pub(crate) fn fix_unnecessary_list_comprehension_set(
-    locator: &Locator,
-    stylist: &Stylist,
+    checker: &Checker,
     expr: &rustpython_parser::ast::Expr,
 ) -> Result<Edit> {
+    let locator = checker.locator;
+    let stylist = checker.stylist;
     // Expr(Call(ListComp)))) ->
     // Expr(SetComp)))
     let module_text = locator.slice(expr.range());
@@ -229,16 +231,21 @@ pub(crate) fn fix_unnecessary_list_comprehension_set(
     };
     tree.codegen(&mut state);
 
-    Ok(Edit::range_replacement(state.to_string(), expr.range()))
+    Ok(Edit::range_replacement(
+        wrap_code_in_spaces(&state, checker.ctx.in_f_string()),
+        expr.range(),
+    ))
 }
 
 /// (C404) Convert `dict([(i, i) for i in range(3)])` to `{i: i for i in
 /// range(3)}`.
 pub(crate) fn fix_unnecessary_list_comprehension_dict(
-    locator: &Locator,
-    stylist: &Stylist,
+    checker: &Checker,
     expr: &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)?;
@@ -254,16 +261,15 @@ pub(crate) fn fix_unnecessary_list_comprehension_dict(
     };
 
     let [Element::Simple {
-            value: key,
-            comma: Some(comma),
+            value: key, ..
         }, Element::Simple { value, .. }] = &tuple.elements[..] else { bail!("Expected tuple with two elements"); };
 
     body.value = Expression::DictComp(Box::new(DictComp {
         key: Box::new(key.clone()),
         value: Box::new(value.clone()),
         for_in: list_comp.for_in.clone(),
-        whitespace_before_colon: comma.whitespace_before.clone(),
-        whitespace_after_colon: comma.whitespace_after.clone(),
+        whitespace_before_colon: ParenthesizableWhitespace::default(),
+        whitespace_after_colon: ParenthesizableWhitespace::SimpleWhitespace(SimpleWhitespace(" ")),
         lbrace: LeftCurlyBrace {
             whitespace_after: call.whitespace_before_args.clone(),
         },
@@ -281,7 +287,10 @@ pub(crate) fn fix_unnecessary_list_comprehension_dict(
     };
     tree.codegen(&mut state);
 
-    Ok(Edit::range_replacement(state.to_string(), expr.range()))
+    Ok(Edit::range_replacement(
+        wrap_code_in_spaces(&state, checker.ctx.in_f_string()),
+        expr.range(),
+    ))
 }
 
 /// Drop a trailing comma from a list of tuple elements.
@@ -329,10 +338,12 @@ fn drop_trailing_comma<'a>(
 
 /// (C405) Convert `set((1, 2))` to `{1, 2}`.
 pub(crate) fn fix_unnecessary_literal_set(
-    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(Set)))
     let module_text = locator.slice(expr.range());
     let mut tree = match_module(module_text)?;
@@ -371,7 +382,10 @@ pub(crate) fn fix_unnecessary_literal_set(
     };
     tree.codegen(&mut state);
 
-    Ok(Edit::range_replacement(state.to_string(), expr.range()))
+    Ok(Edit::range_replacement(
+        wrap_code_in_spaces(&state, checker.ctx.in_f_string()),
+        expr.range(),
+    ))
 }
 
 /// (C406) Convert `dict([(1, 2)])` to `{1: 2}`.
@@ -445,10 +459,12 @@ pub(crate) fn fix_unnecessary_literal_dict(
 
 /// (C408)
 pub(crate) fn fix_unnecessary_collection_call(
-    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" | "dict")))) -> Expr(List|Tuple|Dict)
     let module_text = locator.slice(expr.range());
     let mut tree = match_module(module_text)?;
@@ -457,6 +473,7 @@ pub(crate) fn fix_unnecessary_collection_call(
     let Expression::Name(name) = &call.func.as_ref() else {
         bail!("Expected Expression::Name");
     };
+    let mut wrap_in_spaces = false;
 
     // Arena allocator used to create formatted strings of sufficient lifetime,
     // below.
@@ -480,6 +497,10 @@ pub(crate) fn fix_unnecessary_collection_call(
             }));
         }
         "dict" => {
+            let in_f_string = checker.ctx.in_f_string();
+            if in_f_string {
+                wrap_in_spaces = true;
+            }
             if call.args.is_empty() {
                 body.value = Expression::Dict(Box::new(Dict {
                     elements: vec![],
@@ -489,16 +510,18 @@ pub(crate) fn fix_unnecessary_collection_call(
                     rpar: vec![],
                 }));
             } else {
+                let quote = checker.f_string_quote_style().unwrap_or(stylist.quote());
+
                 // Quote each argument.
                 for arg in &call.args {
                     let quoted = format!(
                         "{}{}{}",
-                        stylist.quote(),
+                        quote,
                         arg.keyword
                             .as_ref()
                             .expect("Expected dictionary argument to be kwarg")
                             .value,
-                        stylist.quote(),
+                        quote,
                     );
                     arena.push(quoted);
                 }
@@ -552,7 +575,19 @@ pub(crate) fn fix_unnecessary_collection_call(
     };
     tree.codegen(&mut state);
 
-    Ok(Edit::range_replacement(state.to_string(), expr.range()))
+    Ok(Edit::range_replacement(
+        wrap_code_in_spaces(&state, wrap_in_spaces),
+        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()
+    }
 }
 
 /// (C409) Convert `tuple([1, 2])` to `tuple(1, 2)`

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

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

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

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

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

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

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

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