[flake8-comprehensions] improve autofix for C401, C402 and C417

handle autofix when parent is f-string

README.md

@@ -1036,9 +1036,9 @@ For more, see [flake8-comprehensions](https://pypi.org/project/flake8-comprehens
 
 | Code | Name | Message | Fix |
 | ---- | ---- | ------- | --- |
-| C400 | unnecessary-generator-list | Unnecessary generator (rewrite as a `list` comprehension) | 🛠 |
-| C401 | unnecessary-generator-set | Unnecessary generator (rewrite as a `set` comprehension) | 🛠 |
-| C402 | unnecessary-generator-dict | Unnecessary generator (rewrite as a `dict` comprehension) | 🛠 |
+| C400 | [unnecessary-generator-list](https://github.com/charliermarsh/ruff/blob/main/docs/rules/unnecessary-generator-list.md) | Unnecessary generator (rewrite as a `list` comprehension) | 🛠 |
+| C401 | [unnecessary-generator-set](https://github.com/charliermarsh/ruff/blob/main/docs/rules/unnecessary-generator-set.md) | Unnecessary generator (rewrite as a `set` comprehension) | 🛠 |
+| C402 | [unnecessary-generator-dict](https://github.com/charliermarsh/ruff/blob/main/docs/rules/unnecessary-generator-dict.md) | Unnecessary generator (rewrite as a `dict` comprehension) | 🛠 |
 | C403 | unnecessary-list-comprehension-set | Unnecessary `list` comprehension (rewrite as a `set` comprehension) | 🛠 |
 | C404 | unnecessary-list-comprehension-dict | Unnecessary `list` comprehension (rewrite as a `dict` comprehension) | 🛠 |
 | C405 | unnecessary-literal-set | Unnecessary `{obj_type}` literal (rewrite as a `set` literal) | 🛠 |

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

@@ -2,7 +2,9 @@
 x = set(
     x for x in range(3)
 )
-
+y = f'{set(a if a < 6 else 0  for a in range(3))}'
+_ = '{}'.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 set(*args, **kwargs):
     return None

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

@@ -3,3 +3,5 @@
     (x, x) for x in range(3)
 )
 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')

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

@@ -11,6 +11,10 @@
 all(map(lambda v: isinstance(v, dict), nums))
 filter(func, map(lambda v: v, nums))
 
+# When inside f-string, then the fix should be surrounded by whitespace
+_ = f"{set(map(lambda x: x % 2 == 0, nums))}"
+_ = f"{dict(map(lambda v: (v, v**2), nums))}"
+
 # Error, but unfixable.
 # For simple expressions, this could be: `(x if x else 1 for x in nums)`.
 # For more complex expressions, this would differ: `(x + 2 if x else 3 for x in nums)`.

crates/ruff/src/checkers/ast.rs

@@ -2434,12 +2434,22 @@ where
                 }
                 if self.settings.rules.enabled(&Rule::UnnecessaryGeneratorSet) {
                     flake8_comprehensions::rules::unnecessary_generator_set(
-                        self, expr, func, args, keywords,
+                        self,
+                        expr,
+                        self.current_expr_parent().map(Into::into),
+                        func,
+                        args,
+                        keywords,
                     );
                 }
                 if self.settings.rules.enabled(&Rule::UnnecessaryGeneratorDict) {
                     flake8_comprehensions::rules::unnecessary_generator_dict(
-                        self, expr, func, args, keywords,
+                        self,
+                        expr,
+                        self.current_expr_parent().map(Into::into),
+                        func,
+                        args,
+                        keywords,
                     );
                 }
                 if self

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

@@ -79,6 +79,7 @@ pub fn fix_unnecessary_generator_set(
     locator: &Locator,
     stylist: &Stylist,
     expr: &rustpython_parser::ast::Expr,
+    parent: Option<&rustpython_parser::ast::Expr>,
 ) -> Result<Fix> {
     // Expr(Call(GeneratorExp)))) -> Expr(SetComp)))
     let module_text = locator.slice_source_code_range(&Range::from_located(expr));
@@ -113,8 +114,17 @@ pub fn fix_unnecessary_generator_set(
     };
     tree.codegen(&mut state);
 
+    let mut content = state.to_string();
+
+    // If parent is f-string then surround with spaces
+    if let Some(parent_element) = parent {
+        if let &rustpython_parser::ast::ExprKind::FormattedValue { .. } = &parent_element.node {
+            content = format!(" {content} ");
+        }
+    }
+
     Ok(Fix::replacement(
-        state.to_string(),
+        content,
         expr.location,
         expr.end_location.unwrap(),
     ))
@@ -126,6 +136,7 @@ pub fn fix_unnecessary_generator_dict(
     locator: &Locator,
     stylist: &Stylist,
     expr: &rustpython_parser::ast::Expr,
+    parent: Option<&rustpython_parser::ast::Expr>,
 ) -> Result<Fix> {
     let module_text = locator.slice_source_code_range(&Range::from_located(expr));
     let mut tree = match_module(module_text)?;
@@ -176,8 +187,17 @@ pub fn fix_unnecessary_generator_dict(
     };
     tree.codegen(&mut state);
 
+    let mut content = state.to_string();
+
+    // If parent is f-string then surround with spaces
+    if let Some(parent_element) = parent {
+        if let &rustpython_parser::ast::ExprKind::FormattedValue { .. } = &parent_element.node {
+            content = format!(" {content} ");
+        }
+    }
+
     Ok(Fix::replacement(
-        state.to_string(),
+        content,
         expr.location,
         expr.end_location.unwrap(),
     ))
@@ -921,6 +941,7 @@ pub fn fix_unnecessary_map(
     locator: &Locator,
     stylist: &Stylist,
     expr: &rustpython_parser::ast::Expr,
+    parent: Option<&rustpython_parser::ast::Expr>,
     kind: &str,
 ) -> Result<Fix> {
     let module_text = locator.slice_source_code_range(&Range::from_located(expr));
@@ -1061,8 +1082,21 @@ pub fn fix_unnecessary_map(
         };
         tree.codegen(&mut state);
 
+        let mut content = state.to_string();
+
+        // If parent is f-string then surround with spaces
+        if kind == "set" || kind == "dict" {
+            if let Some(parent_element) = parent {
+                if let &rustpython_parser::ast::ExprKind::FormattedValue { .. } =
+                    &parent_element.node
+                {
+                    content = format!(" {content} ");
+                }
+            }
+        }
+
         Ok(Fix::replacement(
-            state.to_string(),
+            content,
             expr.location,
             expr.end_location.unwrap(),
         ))

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

@@ -10,6 +10,23 @@ use crate::rules::flake8_comprehensions::fixes;
 use crate::violation::AlwaysAutofixableViolation;
 
 define_violation!(
+    /// ## What it does
+    /// Checks for unnecessary generator that can be rewritten as `dict` comprehension.
+    ///
+    /// ## Why is this bad?
+    /// It is unnecessary to use `dict` around a generator expression, since there are
+    /// equivalent comprehensions for these types.
+    ///
+    /// ## Examples
+    /// ```python
+    /// dict((x, f(x)) for x in foo)
+    /// ```
+    ///
+    /// Use instead:
+    /// ```python
+    /// {x: f(x) for x in foo}
+    /// ```
+    ///
     pub struct UnnecessaryGeneratorDict;
 );
 impl AlwaysAutofixableViolation for UnnecessaryGeneratorDict {
@@ -27,6 +44,7 @@ impl AlwaysAutofixableViolation for UnnecessaryGeneratorDict {
 pub fn unnecessary_generator_dict(
     checker: &mut Checker,
     expr: &Expr,
+    parent: Option<&Expr>,
     func: &Expr,
     args: &[Expr],
     keywords: &[Keyword],
@@ -44,6 +62,7 @@ pub fn unnecessary_generator_dict(
                         checker.locator,
                         checker.stylist,
                         expr,
+                        parent,
                     ) {
                         Ok(fix) => {
                             diagnostic.amend(fix);

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

@@ -10,6 +10,23 @@ use crate::rules::flake8_comprehensions::fixes;
 use crate::violation::AlwaysAutofixableViolation;
 
 define_violation!(
+    /// ## What it does
+    /// Checks for unnecessary generator that can be rewritten as `list` comprehension.
+    ///
+    /// ## Why is this bad?
+    /// It is unnecessary to use `list` around a generator expression, since there are
+    /// equivalent comprehensions for these types.
+    ///
+    /// ## Examples
+    /// ```python
+    /// list(f(x) for x in foo)
+    /// ```
+    ///
+    /// Use instead:
+    /// ```python
+    /// [f(x) for x in foo]
+    /// ```
+    ///
     pub struct UnnecessaryGeneratorList;
 );
 impl AlwaysAutofixableViolation for UnnecessaryGeneratorList {

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

@@ -10,6 +10,23 @@ use crate::rules::flake8_comprehensions::fixes;
 use crate::violation::AlwaysAutofixableViolation;
 
 define_violation!(
+    /// ## What it does
+    /// Checks for unnecessary generator that can be rewritten as `set` comprehension.
+    ///
+    /// ## Why is this bad?
+    /// It is unnecessary to use `set` around a generator expression, since there are
+    /// equivalent comprehensions for these types.
+    ///
+    /// ## Examples
+    /// ```python
+    /// set(f(x) for x in foo)
+    /// ```
+    ///
+    /// Use instead:
+    /// ```python
+    /// {f(x) for x in foo}
+    /// ```
+    ///
     pub struct UnnecessaryGeneratorSet;
 );
 impl AlwaysAutofixableViolation for UnnecessaryGeneratorSet {
@@ -27,6 +44,7 @@ impl AlwaysAutofixableViolation for UnnecessaryGeneratorSet {
 pub fn unnecessary_generator_set(
     checker: &mut Checker,
     expr: &Expr,
+    parent: Option<&Expr>,
     func: &Expr,
     args: &[Expr],
     keywords: &[Keyword],
@@ -40,7 +58,12 @@ pub fn unnecessary_generator_set(
     if let ExprKind::GeneratorExp { .. } = argument {
         let mut diagnostic = Diagnostic::new(UnnecessaryGeneratorSet, Range::from_located(expr));
         if checker.patch(diagnostic.kind.rule()) {
-            match fixes::fix_unnecessary_generator_set(checker.locator, checker.stylist, expr) {
+            match fixes::fix_unnecessary_generator_set(
+                checker.locator,
+                checker.stylist,
+                expr,
+                parent,
+            ) {
                 Ok(fix) => {
                     diagnostic.amend(fix);
                 }

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

@@ -110,6 +110,7 @@ pub fn unnecessary_map(
                         checker.locator,
                         checker.stylist,
                         expr,
+                        parent,
                         "generator",
                     ) {
                         Ok(fix) => {
@@ -141,6 +142,7 @@ pub fn unnecessary_map(
                                 checker.locator,
                                 checker.stylist,
                                 expr,
+                                parent,
                                 id,
                             ) {
                                 Ok(fix) => {
@@ -173,6 +175,7 @@ pub fn unnecessary_map(
                                     checker.locator,
                                     checker.stylist,
                                     expr,
+                                    parent,
                                     id,
                                 ) {
                                     Ok(fix) => {

crates/ruff/src/rules/flake8_comprehensions/snapshots/ruff__rules__flake8_comprehensions__tests__C401_C401.py.snap

@@ -1,5 +1,5 @@
 ---
-source: src/rules/flake8_comprehensions/mod.rs
+source: crates/ruff/src/rules/flake8_comprehensions/mod.rs
 expression: diagnostics
 ---
 - kind:
@@ -40,4 +40,58 @@ expression: diagnostics
       row: 4
       column: 1
   parent: ~
+- kind:
+    UnnecessaryGeneratorSet: ~
+  location:
+    row: 5
+    column: 7
+  end_location:
+    row: 5
+    column: 48
+  fix:
+    content:
+      - " {a if a < 6 else 0  for a in range(3)} "
+    location:
+      row: 5
+      column: 7
+    end_location:
+      row: 5
+      column: 48
+  parent: ~
+- kind:
+    UnnecessaryGeneratorSet: ~
+  location:
+    row: 6
+    column: 16
+  end_location:
+    row: 6
+    column: 57
+  fix:
+    content:
+      - "{a if a < 6 else 0  for a in range(3)}"
+    location:
+      row: 6
+      column: 16
+    end_location:
+      row: 6
+      column: 57
+  parent: ~
+- kind:
+    UnnecessaryGeneratorSet: ~
+  location:
+    row: 7
+    column: 15
+  end_location:
+    row: 7
+    column: 39
+  fix:
+    content:
+      - " {a for a in range(3)} "
+    location:
+      row: 7
+      column: 15
+    end_location:
+      row: 7
+      column: 39
+  parent: ~
 

crates/ruff/src/rules/flake8_comprehensions/snapshots/ruff__rules__flake8_comprehensions__tests__C402_C402.py.snap

@@ -1,5 +1,5 @@
 ---
-source: src/rules/flake8_comprehensions/mod.rs
+source: crates/ruff/src/rules/flake8_comprehensions/mod.rs
 expression: diagnostics
 ---
 - kind:
@@ -40,4 +40,40 @@ expression: diagnostics
       row: 4
       column: 1
   parent: ~
+- kind:
+    UnnecessaryGeneratorDict: ~
+  location:
+    row: 6
+    column: 7
+  end_location:
+    row: 6
+    column: 37
+  fix:
+    content:
+      - " {x: x for x in range(3)} "
+    location:
+      row: 6
+      column: 7
+    end_location:
+      row: 6
+      column: 37
+  parent: ~
+- kind:
+    UnnecessaryGeneratorDict: ~
+  location:
+    row: 7
+    column: 15
+  end_location:
+    row: 7
+    column: 45
+  fix:
+    content:
+      - " {x: x for x in range(3)} "
+    location:
+      row: 7
+      column: 15
+    end_location:
+      row: 7
+      column: 45
+  parent: ~