[flake8-comprehensions] Handle template strings for comprehension fixes (astral-sh#18710)

Essentially this PR ensures that when we do fixes like this:

```diff
- t"{set(f(x) for x in foo)}"
+ t"{ {f(x) for x in foo} }"
```
we are correctly adding whitespace around the braces. 

This logic is already in place for f-strings and just needed to be
generalized to interpolated strings.

Author: dylwil3

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

@@ -36,9 +36,19 @@ def f(x):
     # some more
 )
 
+# t-strings
+print(t"Hello {set(f(a) for a in 'abc')} World")
+print(t"Hello { set(f(a) for a in 'abc') } World")
+small_nums = t"{set(a if a < 6 else 0 for a in range(3))}"
+print(t"Hello {set(a for a in range(3))} World")
+print(t"{set(a for a in 'abc') - set(a for a in 'ab')}")
+print(t"{ set(a for a in 'abc') - set(a for a in 'ab') }")
+
+
 # Not built-in set.
 def set(*args, **kwargs):
     return None
 
 set(2 * x for x in range(3))
 set(x for x in range(3))
+

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

@@ -34,4 +34,16 @@ def f(x):
  ))))
 
 # Test trailing comma case
-s = set([x for x in range(3)],)
+s = set([x for x in range(3)],)
+
+s = t"{set([x for x in 'ab'])}"
+s = t'{set([x for x in "ab"])}'
+
+def f(x):
+    return x
+
+s = t"{set([f(x) for x in 'ab'])}"
+
+s = t"{ set([x for x in 'ab']) | set([x for x in 'ab']) }"
+s = t"{set([x for x in 'ab']) | set([x for x in 'ab'])}"
+

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

@@ -24,3 +24,12 @@
 f"{ set(['a', 'b']) - set(['a']) }"
 f"a {set(['a', 'b']) - set(['a'])} b"
 f"a { set(['a', 'b']) - set(['a']) } b"
+
+t"{set([1,2,3])}"
+t"{set(['a', 'b'])}"
+t'{set(["a", "b"])}'
+
+t"{set(['a', 'b']) - set(['a'])}"
+t"{ set(['a', 'b']) - set(['a']) }"
+t"a {set(['a', 'b']) - set(['a'])} b"
+t"a { set(['a', 'b']) - set(['a']) } b"

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

@@ -27,3 +27,13 @@ def list():
 
 tuple(  # comment
 )
+
+t"{dict(x='y')}"
+t'{dict(x="y")}'
+t"{dict()}"
+t"a {dict()} b"
+
+t"{dict(x='y') | dict(y='z')}"
+t"{ dict(x='y') | dict(y='z') }"
+t"a {dict(x='y') | dict(y='z')} b"
+t"a { dict(x='y') | dict(y='z') } b"

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

@@ -70,3 +70,8 @@ def func(arg1: int, arg2: int = 4):
 list(map(lambda x, y: x, [(1, 2), (3, 4)]))
 list(map(lambda: 1, "xyz"))
 list(map(lambda x, y: x, [(1, 2), (3, 4)]))
+
+# When inside t-string, then the fix should be surrounded by whitespace
+_ = t"{set(map(lambda x: x % 2 == 0, nums))}"
+_ = t"{dict(map(lambda v: (v, v**2), nums))}"
+

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

@@ -37,8 +37,8 @@ use ruff_python_ast::str::Quote;
 use ruff_python_ast::visitor::{Visitor, walk_except_handler, walk_pattern};
 use ruff_python_ast::{
     self as ast, AnyParameterRef, ArgOrKeyword, Comprehension, ElifElseClause, ExceptHandler, Expr,
-    ExprContext, InterpolatedStringElement, Keyword, MatchCase, ModModule, Parameter, Parameters,
-    Pattern, PythonVersion, Stmt, Suite, UnaryOp,
+    ExprContext, ExprFString, ExprTString, InterpolatedStringElement, Keyword, MatchCase,
+    ModModule, Parameter, Parameters, Pattern, PythonVersion, Stmt, Suite, UnaryOp,
 };
 use ruff_python_ast::{PySourceType, helpers, str, visitor};
 use ruff_python_codegen::{Generator, Stylist};
@@ -323,7 +323,8 @@ impl<'a> Checker<'a> {
     /// Return the preferred quote for a generated `StringLiteral` node, given where we are in the
     /// AST.
     fn preferred_quote(&self) -> Quote {
-        self.f_string_quote_style().unwrap_or(self.stylist.quote())
+        self.interpolated_string_quote_style()
+            .unwrap_or(self.stylist.quote())
     }
 
     /// Return the default string flags a generated `StringLiteral` node should use, given where we
@@ -345,21 +346,27 @@ impl<'a> Checker<'a> {
         ast::FStringFlags::empty().with_quote_style(self.preferred_quote())
     }
 
-    /// Returns the appropriate quoting for f-string by reversing the one used outside of
-    /// the f-string.
+    /// Returns the appropriate quoting for interpolated strings by reversing the one used outside of
+    /// the interpolated 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> {
-        if !self.semantic.in_f_string() {
+    /// If the current expression in the context is not an interpolated string, returns ``None``.
+    pub(crate) fn interpolated_string_quote_style(&self) -> Option<Quote> {
+        if !self.semantic.in_interpolated_string() {
             return None;
         }
 
-        // Find the quote character used to start the containing f-string.
-        let ast::ExprFString { value, .. } = self
-            .semantic
+        // Find the quote character used to start the containing interpolated string.
+        self.semantic
             .current_expressions()
-            .find_map(|expr| expr.as_f_string_expr())?;
-        Some(value.iter().next()?.quote_style().opposite())
+            .find_map(|expr| match expr {
+                Expr::FString(ExprFString { value, .. }) => {
+                    Some(value.iter().next()?.quote_style().opposite())
+                }
+                Expr::TString(ExprTString { value, .. }) => {
+                    Some(value.iter().next()?.quote_style().opposite())
+                }
+                _ => None,
+            })
     }
 
     /// Returns the [`SourceRow`] for the given offset.

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

@@ -236,7 +236,9 @@ pub(crate) fn fix_unnecessary_collection_call(
     // below.
     let mut arena: Vec<String> = vec![];
 
-    let quote = checker.f_string_quote_style().unwrap_or(stylist.quote());
+    let quote = checker
+        .interpolated_string_quote_style()
+        .unwrap_or(stylist.quote());
 
     // Quote each argument.
     for arg in &call.args {
@@ -317,7 +319,7 @@ pub(crate) fn pad_expression(
     locator: &Locator,
     semantic: &SemanticModel,
 ) -> String {
-    if !semantic.in_f_string() {
+    if !semantic.in_interpolated_string() {
         return content;
     }
 
@@ -349,7 +351,7 @@ pub(crate) fn pad_start(
     locator: &Locator,
     semantic: &SemanticModel,
 ) -> String {
-    if !semantic.in_f_string() {
+    if !semantic.in_interpolated_string() {
         return content.into();
     }
 
@@ -370,7 +372,7 @@ pub(crate) fn pad_end(
     locator: &Locator,
     semantic: &SemanticModel,
 ) -> String {
-    if !semantic.in_f_string() {
+    if !semantic.in_interpolated_string() {
         return content.into();
     }
 
@@ -798,10 +800,10 @@ pub(crate) fn fix_unnecessary_map(
 
     let mut content = tree.codegen_stylist(stylist);
 
-    // If the expression is embedded in an f-string, surround it with spaces to avoid
+    // If the expression is embedded in an interpolated string, surround it with spaces to avoid
     // syntax errors.
     if matches!(object_type, ObjectType::Set | ObjectType::Dict) {
-        if parent.is_some_and(Expr::is_f_string_expr) {
+        if parent.is_some_and(|expr| expr.is_f_string_expr() || expr.is_t_string_expr()) {
             content = format!(" {content} ");
         }
     }

crates/ruff_linter/src/rules/flake8_comprehensions/snapshots/ruff_linter__rules__flake8_comprehensions__tests__C401_C401.py.snap

@@ -346,7 +346,7 @@ C401.py:32:1: C401 [*] Unnecessary generator (rewrite as a set comprehension)
 37 | | )
    | |_^ C401
 38 |
-39 |   # Not built-in set.
+39 |   # t-strings
    |
    = help: Rewrite as a set comprehension
 
@@ -364,5 +364,166 @@ C401.py:32:1: C401 [*] Unnecessary generator (rewrite as a set comprehension)
 37    |-)
    35 |+    }
 38 36 | 
-39 37 | # Not built-in set.
-40 38 | def set(*args, **kwargs):
+39 37 | # t-strings
+40 38 | print(t"Hello {set(f(a) for a in 'abc')} World")
+
+C401.py:40:16: C401 [*] Unnecessary generator (rewrite as a set comprehension)
+   |
+39 | # t-strings
+40 | print(t"Hello {set(f(a) for a in 'abc')} World")
+   |                ^^^^^^^^^^^^^^^^^^^^^^^^ C401
+41 | print(t"Hello { set(f(a) for a in 'abc') } World")
+42 | small_nums = t"{set(a if a < 6 else 0 for a in range(3))}"
+   |
+   = help: Rewrite as a set comprehension
+
+ℹ Unsafe fix
+37 37 | )
+38 38 | 
+39 39 | # t-strings
+40    |-print(t"Hello {set(f(a) for a in 'abc')} World")
+   40 |+print(t"Hello { {f(a) for a in 'abc'} } World")
+41 41 | print(t"Hello { set(f(a) for a in 'abc') } World")
+42 42 | small_nums = t"{set(a if a < 6 else 0 for a in range(3))}"
+43 43 | print(t"Hello {set(a for a in range(3))} World")
+
+C401.py:41:17: C401 [*] Unnecessary generator (rewrite as a set comprehension)
+   |
+39 | # t-strings
+40 | print(t"Hello {set(f(a) for a in 'abc')} World")
+41 | print(t"Hello { set(f(a) for a in 'abc') } World")
+   |                 ^^^^^^^^^^^^^^^^^^^^^^^^ C401
+42 | small_nums = t"{set(a if a < 6 else 0 for a in range(3))}"
+43 | print(t"Hello {set(a for a in range(3))} World")
+   |
+   = help: Rewrite as a set comprehension
+
+ℹ Unsafe fix
+38 38 | 
+39 39 | # t-strings
+40 40 | print(t"Hello {set(f(a) for a in 'abc')} World")
+41    |-print(t"Hello { set(f(a) for a in 'abc') } World")
+   41 |+print(t"Hello { {f(a) for a in 'abc'} } World")
+42 42 | small_nums = t"{set(a if a < 6 else 0 for a in range(3))}"
+43 43 | print(t"Hello {set(a for a in range(3))} World")
+44 44 | print(t"{set(a for a in 'abc') - set(a for a in 'ab')}")
+
+C401.py:42:17: C401 [*] Unnecessary generator (rewrite as a set comprehension)
+   |
+40 | print(t"Hello {set(f(a) for a in 'abc')} World")
+41 | print(t"Hello { set(f(a) for a in 'abc') } World")
+42 | small_nums = t"{set(a if a < 6 else 0 for a in range(3))}"
+   |                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ C401
+43 | print(t"Hello {set(a for a in range(3))} World")
+44 | print(t"{set(a for a in 'abc') - set(a for a in 'ab')}")
+   |
+   = help: Rewrite as a set comprehension
+
+ℹ Unsafe fix
+39 39 | # t-strings
+40 40 | print(t"Hello {set(f(a) for a in 'abc')} World")
+41 41 | print(t"Hello { set(f(a) for a in 'abc') } World")
+42    |-small_nums = t"{set(a if a < 6 else 0 for a in range(3))}"
+   42 |+small_nums = t"{ {a if a < 6 else 0 for a in range(3)} }"
+43 43 | print(t"Hello {set(a for a in range(3))} World")
+44 44 | print(t"{set(a for a in 'abc') - set(a for a in 'ab')}")
+45 45 | print(t"{ set(a for a in 'abc') - set(a for a in 'ab') }")
+
+C401.py:43:16: C401 [*] Unnecessary generator (rewrite using `set()`)
+   |
+41 | print(t"Hello { set(f(a) for a in 'abc') } World")
+42 | small_nums = t"{set(a if a < 6 else 0 for a in range(3))}"
+43 | print(t"Hello {set(a for a in range(3))} World")
+   |                ^^^^^^^^^^^^^^^^^^^^^^^^ C401
+44 | print(t"{set(a for a in 'abc') - set(a for a in 'ab')}")
+45 | print(t"{ set(a for a in 'abc') - set(a for a in 'ab') }")
+   |
+   = help: Rewrite using `set()`
+
+ℹ Unsafe fix
+40 40 | print(t"Hello {set(f(a) for a in 'abc')} World")
+41 41 | print(t"Hello { set(f(a) for a in 'abc') } World")
+42 42 | small_nums = t"{set(a if a < 6 else 0 for a in range(3))}"
+43    |-print(t"Hello {set(a for a in range(3))} World")
+   43 |+print(t"Hello {set(range(3))} World")
+44 44 | print(t"{set(a for a in 'abc') - set(a for a in 'ab')}")
+45 45 | print(t"{ set(a for a in 'abc') - set(a for a in 'ab') }")
+46 46 | 
+
+C401.py:44:10: C401 [*] Unnecessary generator (rewrite using `set()`)
+   |
+42 | small_nums = t"{set(a if a < 6 else 0 for a in range(3))}"
+43 | print(t"Hello {set(a for a in range(3))} World")
+44 | print(t"{set(a for a in 'abc') - set(a for a in 'ab')}")
+   |          ^^^^^^^^^^^^^^^^^^^^^ C401
+45 | print(t"{ set(a for a in 'abc') - set(a for a in 'ab') }")
+   |
+   = help: Rewrite using `set()`
+
+ℹ Unsafe fix
+41 41 | print(t"Hello { set(f(a) for a in 'abc') } World")
+42 42 | small_nums = t"{set(a if a < 6 else 0 for a in range(3))}"
+43 43 | print(t"Hello {set(a for a in range(3))} World")
+44    |-print(t"{set(a for a in 'abc') - set(a for a in 'ab')}")
+   44 |+print(t"{set('abc') - set(a for a in 'ab')}")
+45 45 | print(t"{ set(a for a in 'abc') - set(a for a in 'ab') }")
+46 46 | 
+47 47 | 
+
+C401.py:44:34: C401 [*] Unnecessary generator (rewrite using `set()`)
+   |
+42 | small_nums = t"{set(a if a < 6 else 0 for a in range(3))}"
+43 | print(t"Hello {set(a for a in range(3))} World")
+44 | print(t"{set(a for a in 'abc') - set(a for a in 'ab')}")
+   |                                  ^^^^^^^^^^^^^^^^^^^^ C401
+45 | print(t"{ set(a for a in 'abc') - set(a for a in 'ab') }")
+   |
+   = help: Rewrite using `set()`
+
+ℹ Unsafe fix
+41 41 | print(t"Hello { set(f(a) for a in 'abc') } World")
+42 42 | small_nums = t"{set(a if a < 6 else 0 for a in range(3))}"
+43 43 | print(t"Hello {set(a for a in range(3))} World")
+44    |-print(t"{set(a for a in 'abc') - set(a for a in 'ab')}")
+   44 |+print(t"{set(a for a in 'abc') - set('ab')}")
+45 45 | print(t"{ set(a for a in 'abc') - set(a for a in 'ab') }")
+46 46 | 
+47 47 | 
+
+C401.py:45:11: C401 [*] Unnecessary generator (rewrite using `set()`)
+   |
+43 | print(t"Hello {set(a for a in range(3))} World")
+44 | print(t"{set(a for a in 'abc') - set(a for a in 'ab')}")
+45 | print(t"{ set(a for a in 'abc') - set(a for a in 'ab') }")
+   |           ^^^^^^^^^^^^^^^^^^^^^ C401
+   |
+   = help: Rewrite using `set()`
+
+ℹ Unsafe fix
+42 42 | small_nums = t"{set(a if a < 6 else 0 for a in range(3))}"
+43 43 | print(t"Hello {set(a for a in range(3))} World")
+44 44 | print(t"{set(a for a in 'abc') - set(a for a in 'ab')}")
+45    |-print(t"{ set(a for a in 'abc') - set(a for a in 'ab') }")
+   45 |+print(t"{ set('abc') - set(a for a in 'ab') }")
+46 46 | 
+47 47 | 
+48 48 | # Not built-in set.
+
+C401.py:45:35: C401 [*] Unnecessary generator (rewrite using `set()`)
+   |
+43 | print(t"Hello {set(a for a in range(3))} World")
+44 | print(t"{set(a for a in 'abc') - set(a for a in 'ab')}")
+45 | print(t"{ set(a for a in 'abc') - set(a for a in 'ab') }")
+   |                                   ^^^^^^^^^^^^^^^^^^^^ C401
+   |
+   = help: Rewrite using `set()`
+
+ℹ Unsafe fix
+42 42 | small_nums = t"{set(a if a < 6 else 0 for a in range(3))}"
+43 43 | print(t"Hello {set(a for a in range(3))} World")
+44 44 | print(t"{set(a for a in 'abc') - set(a for a in 'ab')}")
+45    |-print(t"{ set(a for a in 'abc') - set(a for a in 'ab') }")
+   45 |+print(t"{ set(a for a in 'abc') - set('ab') }")
+46 46 | 
+47 47 | 
+48 48 | # Not built-in set.