From 86a905154f69fba963cdbea874b49231df1bce7c Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Mon, 30 Oct 2023 19:27:15 -0400
Subject: [PATCH] Fix invalid E231 error with f-strings

---
 .../test/fixtures/pycodestyle/E23.py          |  6 ++++++
 .../rules/logical_lines/missing_whitespace.rs | 10 +++++++---
 ...ules__pycodestyle__tests__E231_E23.py.snap | 20 +++++++++++++++++++
 3 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/crates/ruff_linter/resources/test/fixtures/pycodestyle/E23.py b/crates/ruff_linter/resources/test/fixtures/pycodestyle/E23.py
index 4d0089bb451ab9..2d7e70e99d4acc 100644
--- a/crates/ruff_linter/resources/test/fixtures/pycodestyle/E23.py
+++ b/crates/ruff_linter/resources/test/fixtures/pycodestyle/E23.py
@@ -40,5 +40,11 @@ def foo() -> None:
 f"{(lambda x:x)}"
 f"normal{f"{a:.3f}"}normal"
 
+#: Okay
+snapshot.file_uri[len(f's3://{self.s3_bucket_name}/'):]
+
+#: E231
+{len(f's3://{self.s3_bucket_name}/'):1}
+
 #: Okay
 a = (1,
diff --git a/crates/ruff_linter/src/rules/pycodestyle/rules/logical_lines/missing_whitespace.rs b/crates/ruff_linter/src/rules/pycodestyle/rules/logical_lines/missing_whitespace.rs
index 314693446aa868..1d8a9caf78f5c1 100644
--- a/crates/ruff_linter/src/rules/pycodestyle/rules/logical_lines/missing_whitespace.rs
+++ b/crates/ruff_linter/src/rules/pycodestyle/rules/logical_lines/missing_whitespace.rs
@@ -64,14 +64,14 @@ pub(crate) fn missing_whitespace(line: &LogicalLine, context: &mut LogicalLinesC
         match kind {
             TokenKind::FStringStart => fstrings += 1,
             TokenKind::FStringEnd => fstrings = fstrings.saturating_sub(1),
-            TokenKind::Lsqb => {
+            TokenKind::Lsqb if fstrings == 0 => {
                 open_parentheses = open_parentheses.saturating_add(1);
                 prev_lsqb = token.start();
             }
-            TokenKind::Rsqb => {
+            TokenKind::Rsqb if fstrings == 0 => {
                 open_parentheses = open_parentheses.saturating_sub(1);
             }
-            TokenKind::Lbrace => {
+            TokenKind::Lbrace if fstrings == 0 => {
                 prev_lbrace = token.start();
             }
             TokenKind::Colon if fstrings > 0 => {
@@ -111,6 +111,10 @@ pub(crate) fn missing_whitespace(line: &LogicalLine, context: &mut LogicalLinesC
                         }
                     }
 
+                    println!("open_parentheses: {}", open_parentheses);
+                    println!("prev_lsqb: {:?}", prev_lsqb);
+                    println!("prev_lbrace: {:?}", prev_lbrace);
+
                     let mut diagnostic =
                         Diagnostic::new(MissingWhitespace { token: kind }, token.range());
                     diagnostic.set_fix(Fix::safe_edit(Edit::insertion(
diff --git a/crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E231_E23.py.snap b/crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E231_E23.py.snap
index d0839a77a266fa..e8e3663c82d521 100644
--- a/crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E231_E23.py.snap
+++ b/crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E231_E23.py.snap
@@ -121,4 +121,24 @@ E23.py:33:6: E231 [*] Missing whitespace after ','
 35 35 | # Okay because it's hard to differentiate between the usages of a colon in a f-string
 36 36 | f"{a:=1}"
 
+E23.py:47:37: E231 [*] Missing whitespace after ':'
+   |
+46 | #: E231
+47 | {len(f's3://{self.s3_bucket_name}/'):1}
+   |                                     ^ E231
+48 | 
+49 | #: Okay
+   |
+   = help: Add missing whitespace
+
+ℹ Fix
+44 44 | snapshot.file_uri[len(f's3://{self.s3_bucket_name}/'):]
+45 45 | 
+46 46 | #: E231
+47    |-{len(f's3://{self.s3_bucket_name}/'):1}
+   47 |+{len(f's3://{self.s3_bucket_name}/'): 1}
+48 48 | 
+49 49 | #: Okay
+50 50 | a = (1,
+