Update logical line rules for the new f-string tokens (#7690)

## Summary

This PR updates the logical line rules for the new f-string tokens.

Specifically, it tries to ignore the rules for which it might be hard to
differentiate between different usages inside f-strings. For some
tokens, the whitespace addition or removal could change the semantics so
they're ignored for every scenario inside f-string.

- Ignore `E225` for the `=` operator. For example, `f"{x=}"` we don't
want to add whitespace as the output will then be different than what
the user indended.
- Ignore `E231` for the `:` operator. For example, `f"{x:.3f}"` we don't
want to add whitespace as the output will be different:
```
In [1]: a = 100

In [2]: f'{a: .3f}'
Out[2]: ' 100.000'

In [3]: f'{a:.3f}'
Out[3]: '100.000'
```

- Ignore `E201`, `E202` for `{` and `}` tokens. The double braces are
used as an escape mechanism to include the raw character (`f"{{foo}}"`)
but if a dictionary is being used then the whitespace is required (`f"{
{'a': 1} }"`).

## Test Plan

Add new test cases for f-strings along with nested f-strings and update
the snapshots.

crates/ruff_linter/resources/test/fixtures/pycodestyle/E20.py

@@ -84,3 +84,8 @@
 x = [  #
     'some value',
 ]
+
+# F-strings
+f"{ {'a': 1} }"
+f"{[ { {'a': 1} } ]}"
+f"normal { {f"{ { [1, 2] } }" } } normal"

crates/ruff_linter/resources/test/fixtures/pycodestyle/E23.py

@@ -29,5 +29,16 @@ def foo() -> None:
     'tag_smalldata':[('byte_count_mdtype', 'u4'), ('data', 'S4')],
 }
 
+# E231
+f"{(a,b)}"
+
+# Okay because it's hard to differentiate between the usages of a colon in a f-string
+f"{a:=1}"
+f"{ {'a':1} }"
+f"{a:.3f}"
+f"{(a:=1)}"
+f"{(lambda x:x)}"
+f"normal{f"{a:.3f}"}normal"
+
 #: Okay
 a = (1,

crates/ruff_linter/resources/test/fixtures/pycodestyle/E25.py

@@ -48,3 +48,9 @@ def add(a: int=0, b: int =0, c: int= 0) -> int:
 #: Okay
 def add(a: int = _default(name='f')):
     return a
+
+# F-strings
+f"{a=}"
+f"{a:=1}"
+f"{foo(a=1)}"
+f"normal {f"{a=}"} normal"

crates/ruff_linter/src/rules/pycodestyle/rules/logical_lines/extraneous_whitespace.rs

@@ -134,12 +134,27 @@ pub(crate) fn extraneous_whitespace(
     autofix_before_punctuation: bool,
 ) {
     let mut prev_token = None;
+    let mut fstrings = 0u32;
 
     for token in line.tokens() {
         let kind = token.kind();
+        match kind {
+            TokenKind::FStringStart => fstrings += 1,
+            TokenKind::FStringEnd => fstrings = fstrings.saturating_sub(1),
+            _ => {}
+        }
         if let Some(symbol) = BracketOrPunctuation::from_kind(kind) {
+            // Whitespace before "{" or after "}" might be required in f-strings.
+            // For example,
+            //
+            // ```python
+            // f"{ {'a': 1} }"
+            // ```
+            //
+            // Here, `{{` / `}} would be interpreted as a single raw `{` / `}`
+            // character.
             match symbol {
-                BracketOrPunctuation::OpenBracket(symbol) => {
+                BracketOrPunctuation::OpenBracket(symbol) if symbol != '{' || fstrings == 0 => {
                     let (trailing, trailing_len) = line.trailing_whitespace(token);
                     if !matches!(trailing, Whitespace::None) {
                         let mut diagnostic = Diagnostic::new(
@@ -153,7 +168,7 @@ pub(crate) fn extraneous_whitespace(
                         context.push_diagnostic(diagnostic);
                     }
                 }
-                BracketOrPunctuation::CloseBracket(symbol) => {
+                BracketOrPunctuation::CloseBracket(symbol) if symbol != '}' || fstrings == 0 => {
                     if !matches!(prev_token, Some(TokenKind::Comma)) {
                         if let (Whitespace::Single | Whitespace::Many | Whitespace::Tab, offset) =
                             line.leading_whitespace(token)
@@ -189,6 +204,7 @@ pub(crate) fn extraneous_whitespace(
                         }
                     }
                 }
+                _ => {}
             }
         }
 

crates/ruff_linter/src/rules/pycodestyle/rules/logical_lines/missing_whitespace.rs

@@ -59,13 +59,16 @@ pub(crate) fn missing_whitespace(
     context: &mut LogicalLinesContext,
 ) {
     let mut open_parentheses = 0u32;
+    let mut fstrings = 0u32;
     let mut prev_lsqb = TextSize::default();
     let mut prev_lbrace = TextSize::default();
     let mut iter = line.tokens().iter().peekable();
 
     while let Some(token) = iter.next() {
         let kind = token.kind();
         match kind {
+            TokenKind::FStringStart => fstrings += 1,
+            TokenKind::FStringEnd => fstrings = fstrings.saturating_sub(1),
             TokenKind::Lsqb => {
                 open_parentheses = open_parentheses.saturating_add(1);
                 prev_lsqb = token.start();
@@ -76,6 +79,17 @@ pub(crate) fn missing_whitespace(
             TokenKind::Lbrace => {
                 prev_lbrace = token.start();
             }
+            TokenKind::Colon if fstrings > 0 => {
+                // Colon in f-string, no space required. This will yield false
+                // negatives for cases like the following as it's hard to
+                // differentiate between the usage of a colon in a f-string.
+                //
+                // ```python
+                // f'{ {'x':1} }'
+                // f'{(lambda x:x)}'
+                // ```
+                continue;
+            }
 
             TokenKind::Comma | TokenKind::Semi | TokenKind::Colon => {
                 let after = line.text_after(token);

crates/ruff_linter/src/rules/pycodestyle/rules/logical_lines/missing_whitespace_around_operator.rs

@@ -141,6 +141,7 @@ pub(crate) fn missing_whitespace_around_operator(
         prev_token.kind(),
         TokenKind::Lpar | TokenKind::Lambda
     ));
+    let mut fstrings = u32::from(matches!(prev_token.kind(), TokenKind::FStringStart));
 
     while let Some(token) = tokens.next() {
         let kind = token.kind();
@@ -150,13 +151,15 @@ pub(crate) fn missing_whitespace_around_operator(
         }
 
         match kind {
+            TokenKind::FStringStart => fstrings += 1,
+            TokenKind::FStringEnd => fstrings = fstrings.saturating_sub(1),
             TokenKind::Lpar | TokenKind::Lambda => parens += 1,
             TokenKind::Rpar => parens = parens.saturating_sub(1),
             _ => {}
         };
 
-        let needs_space = if kind == TokenKind::Equal && parens > 0 {
-            // Allow keyword args or defaults: foo(bar=None).
+        let needs_space = if kind == TokenKind::Equal && (parens > 0 || fstrings > 0) {
+            // Allow keyword args, defaults: foo(bar=None) and f-strings: f'{foo=}'
             NeedsSpace::No
         } else if kind == TokenKind::Slash {
             // Tolerate the "/" operator in function definition

crates/ruff_linter/src/rules/pycodestyle/rules/logical_lines/whitespace_around_named_parameter_equals.rs

@@ -26,7 +26,7 @@ use crate::rules::pycodestyle::rules::logical_lines::{LogicalLine, LogicalLineTo
 ///
 /// Use instead:
 /// ```python
-/// def add(a = 0) -> int:
+/// def add(a=0) -> int:
 ///     return a + 1
 /// ```
 ///

crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E201_E20.py.snap

@@ -144,4 +144,22 @@ E20.py:81:6: E201 [*] Whitespace after '['
 83 83 | #: Okay
 84 84 | x = [  #
 
+E20.py:90:5: E201 [*] Whitespace after '['
+   |
+88 | # F-strings
+89 | f"{ {'a': 1} }"
+90 | f"{[ { {'a': 1} } ]}"
+   |     ^ E201
+91 | f"normal { {f"{ { [1, 2] } }" } } normal"
+   |
+   = help: Remove whitespace before '['
+
+ℹ Fix
+87 87 | 
+88 88 | # F-strings
+89 89 | f"{ {'a': 1} }"
+90    |-f"{[ { {'a': 1} } ]}"
+   90 |+f"{[{ {'a': 1} } ]}"
+91 91 | f"normal { {f"{ { [1, 2] } }" } } normal"
+
 

crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E202_E20.py.snap

@@ -126,4 +126,22 @@ E20.py:29:11: E202 [*] Whitespace before ']'
 31 31 | spam(ham[1], {eggs: 2})
 32 32 | 
 
+E20.py:90:18: E202 [*] Whitespace before ']'
+   |
+88 | # F-strings
+89 | f"{ {'a': 1} }"
+90 | f"{[ { {'a': 1} } ]}"
+   |                  ^ E202
+91 | f"normal { {f"{ { [1, 2] } }" } } normal"
+   |
+   = help: Remove whitespace before ']'
+
+ℹ Fix
+87 87 | 
+88 88 | # F-strings
+89 89 | f"{ {'a': 1} }"
+90    |-f"{[ { {'a': 1} } ]}"
+   90 |+f"{[ { {'a': 1} }]}"
+91 91 | f"normal { {f"{ { [1, 2] } }" } } normal"
+
 

crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E231_E23.py.snap

@@ -99,6 +99,26 @@ E23.py:29:20: E231 [*] Missing whitespace after ':'
    29 |+    'tag_smalldata': [('byte_count_mdtype', 'u4'), ('data', 'S4')],
 30 30 | }
 31 31 | 
-32 32 | #: Okay
+32 32 | # E231
+
+E23.py:33:6: E231 [*] Missing whitespace after ','
+   |
+32 | # E231
+33 | f"{(a,b)}"
+   |      ^ E231
+34 | 
+35 | # Okay because it's hard to differentiate between the usages of a colon in a f-string
+   |
+   = help: Added missing whitespace after ','
+
+ℹ Fix
+30 30 | }
+31 31 | 
+32 32 | # E231
+33    |-f"{(a,b)}"
+   33 |+f"{(a, b)}"
+34 34 | 
+35 35 | # Okay because it's hard to differentiate between the usages of a colon in a f-string
+36 36 | f"{a:=1}"