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

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

@@ -1,13 +1,20 @@
 x = set(x for x in range(3))
-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')
+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 f(x):
+    return x
 
-def set(*args, **kwargs):
-    return None
 
+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') }")
 
-set(x for x in range(3))
+# The fix generated for this diagnostic is incorrect, as we add additional space
+# around the set comprehension.
+print(f"{ {set(a for a in 'abc')} }")

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

@@ -5,3 +5,14 @@
 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')
+
+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
+
+print(f'Hello {dict((x,f(x)) for x in "abc")} World')

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

@@ -2,3 +2,14 @@
 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'])}"
+
+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

@@ -1,2 +1,13 @@
 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'])}"
+
+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

@@ -16,3 +16,11 @@
 set(
     [1,]
 )
+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']) }"
+f"a {set(['a', 'b']) - set(['a'])} b"
+f"a { set(['a', 'b']) - set(['a']) } b"

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

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

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

@@ -136,36 +136,36 @@ 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(
-            model: &SemanticModel,
-            locator: &Locator,
-            indexer: &Indexer,
-        ) -> Option<Quote> {
-            if !model.in_f_string() {
-                return None;
-            }
-
-            // Find the quote character used to start the containing f-string.
-            let expr = model.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.semantic_model, 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 model = &self.semantic_model;
+        if !model.in_f_string() {
+            return None;
+        }
+
+        // Find the quote character used to start the containing f-string.
+        let expr = model.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,
+        }
+    }
+
     /// Returns the [`IsolationLevel`] for fixes in the current context.
     ///
     /// The primary use-case for fix isolation is to ensure that we don't delete all statements
@@ -2729,22 +2729,12 @@ where
                 }
                 if self.enabled(Rule::UnnecessaryGeneratorSet) {
                     flake8_comprehensions::rules::unnecessary_generator_set(
-                        self,
-                        expr,
-                        self.semantic_model.expr_parent(),
-                        func,
-                        args,
-                        keywords,
+                        self, expr, func, args, keywords,
                     );
                 }
                 if self.enabled(Rule::UnnecessaryGeneratorDict) {
                     flake8_comprehensions::rules::unnecessary_generator_dict(
-                        self,
-                        expr,
-                        self.semantic_model.expr_parent(),
-                        func,
-                        args,
-                        keywords,
+                        self, expr, func, args, keywords,
                     );
                 }
                 if self.enabled(Rule::UnnecessaryListComprehensionSet) {