Allow whitespace around colon in slices for E203

astral-sh · Nov 13, 2023 · 8816092 · 8816092
1 parent 6a6de53
commit 8816092
Show file tree

Hide file tree

Showing 5 changed files with 229 additions and 10 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/pycodestyle/E20.py b/crates/ruff_linter/resources/test/fixtures/pycodestyle/E20.py
@@ -89,3 +89,49 @@
 f"{ {'a': 1} }"
 f"{[ { {'a': 1} } ]}"
 f"normal { {f"{ { [1, 2] } }" } } normal"
+
+#: Okay
+ham[lower + offset : upper + offset]
+
+#: Okay
+ham[(lower + offset) : upper + offset]
+
+#: E203:1:19
+ham{lower + offset : upper + offset}
+
+#: E203:1:19
+ham[lower + offset  : upper + offset]
+
+#: Okay
+release_lines = history_file_lines[history_file_lines.index('## Unreleased') + 1: -1]
+
+#: Okay
+release_lines = history_file_lines[history_file_lines.index('## Unreleased') + 1 : -1]
+
+#: Okay
+ham[1:9], ham[1:9:3], ham[:9:3], ham[1::3], ham[1:9:]
+ham[lower:upper], ham[lower:upper:], ham[lower::step]
+ham[lower+offset : upper+offset]
+ham[: upper_fn(x) : step_fn(x)], ham[:: step_fn(x)]
+ham[lower + offset : upper + offset]
+
+#: E201:1:5
+ham[ : upper]
+
+#: Okay
+ham[lower + offset :: upper + offset]
+
+#: Okay
+ham[(lower + offset) :: upper + offset]
+
+#: Okay
+ham[lower + offset::upper + offset]
+
+#: E203:1:21
+ham[lower + offset : : upper + offset]
+
+#: E203:1:20
+ham[lower + offset: :upper + offset]
+
+#: E203:1:20
+ham[{lower + offset : upper + offset} : upper + offset]
diff --git a/crates/ruff_linter/src/rules/pycodestyle/rules/logical_lines/extraneous_whitespace.rs b/crates/ruff_linter/src/rules/pycodestyle/rules/logical_lines/extraneous_whitespace.rs
@@ -127,14 +127,28 @@ impl AlwaysFixableViolation for WhitespaceBeforePunctuation {
 
 /// E201, E202, E203
 pub(crate) fn extraneous_whitespace(line: &LogicalLine, context: &mut LogicalLinesContext) {
-    let mut prev_token = None;
     let mut fstrings = 0u32;
+    let mut brackets = vec![];
+    let mut prev_token = None;
+    let mut iter = line.tokens().iter().peekable();
 
-    for token in line.tokens() {
+    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 if fstrings == 0 => {
+                brackets.push(kind);
+            }
+            TokenKind::Rsqb if fstrings == 0 => {
+                brackets.pop();
+            }
+            TokenKind::Lbrace if fstrings == 0 => {
+                brackets.push(kind);
+            }
+            TokenKind::Rbrace if fstrings == 0 => {
+                brackets.pop();
+            }
             _ => {}
         }
         if let Some(symbol) = BracketOrPunctuation::from_kind(kind) {
@@ -177,16 +191,56 @@ pub(crate) fn extraneous_whitespace(line: &LogicalLine, context: &mut LogicalLin
                 }
                 BracketOrPunctuation::Punctuation(symbol) => {
                     if !matches!(prev_token, Some(TokenKind::Comma)) {
+                        let whitespace = line.leading_whitespace(token);
                         if let (Whitespace::Single | Whitespace::Many | Whitespace::Tab, offset) =
-                            line.leading_whitespace(token)
+                            whitespace
                         {
-                            let mut diagnostic = Diagnostic::new(
-                                WhitespaceBeforePunctuation { symbol },
-                                TextRange::at(token.start() - offset, offset),
-                            );
-                            diagnostic
-                                .set_fix(Fix::safe_edit(Edit::range_deletion(diagnostic.range())));
-                            context.push_diagnostic(diagnostic);
+                            // If we're in a slice, and the token is a colon, and it has
+                            // equivalent spacing on both sides, allow it.
+                            if symbol == ':'
+                                && brackets
+                                    .last()
+                                    .is_some_and(|kind| matches!(kind, TokenKind::Lsqb))
+                            {
+                                // If we're in the second half of a double colon, disallow
+                                // any whitespace (e.g., `foo[1: :2]` or `foo[1 : : 2]`).
+                                if matches!(prev_token, Some(TokenKind::Colon)) {
+                                    let mut diagnostic = Diagnostic::new(
+                                        WhitespaceBeforePunctuation { symbol },
+                                        TextRange::at(token.start() - offset, offset),
+                                    );
+                                    diagnostic.set_fix(Fix::safe_edit(Edit::range_deletion(
+                                        diagnostic.range(),
+                                    )));
+                                    context.push_diagnostic(diagnostic);
+                                } else {
+                                    let token = iter
+                                        .peek()
+                                        .filter(|next| matches!(next.kind(), TokenKind::Colon))
+                                        .unwrap_or(&token);
+
+                                    // Allow, e.g., `foo[1:2]` or `foo[1 : 2]` or `foo[1 :: 2]`.
+                                    if line.trailing_whitespace(token) != whitespace {
+                                        let mut diagnostic = Diagnostic::new(
+                                            WhitespaceBeforePunctuation { symbol },
+                                            TextRange::at(token.start() - offset, offset),
+                                        );
+                                        diagnostic.set_fix(Fix::safe_edit(Edit::range_deletion(
+                                            diagnostic.range(),
+                                        )));
+                                        context.push_diagnostic(diagnostic);
+                                    }
+                                }
+                            } else {
+                                let mut diagnostic = Diagnostic::new(
+                                    WhitespaceBeforePunctuation { symbol },
+                                    TextRange::at(token.start() - offset, offset),
+                                );
+                                diagnostic.set_fix(Fix::safe_edit(Edit::range_deletion(
+                                    diagnostic.range(),
+                                )));
+                                context.push_diagnostic(diagnostic);
+                            }
                         }
                     }
                 }

diff --git a/.../src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E201_E20.py.snap b/.../src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E201_E20.py.snap
@@ -161,5 +161,27 @@ E20.py:90:5: E201 [*] Whitespace after '['
 90    |-f"{[ { {'a': 1} } ]}"
    90 |+f"{[{ {'a': 1} } ]}"
 91 91 | f"normal { {f"{ { [1, 2] } }" } } normal"
+92 92 | 
+93 93 | #: Okay
+
+E20.py:119:5: E201 [*] Whitespace after '['
+    |
+118 | #: E201:1:5
+119 | ham[ : upper]
+    |     ^ E201
+120 | 
+121 | #: Okay
+    |
+    = help: Remove whitespace before '['
+
+ℹ Safe fix
+116 116 | ham[lower + offset : upper + offset]
+117 117 | 
+118 118 | #: E201:1:5
+119     |-ham[ : upper]
+    119 |+ham[: upper]
+120 120 | 
+121 121 | #: Okay
+122 122 | ham[lower + offset :: upper + offset]
 
 
diff --git a/.../src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E202_E20.py.snap b/.../src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E202_E20.py.snap
@@ -143,5 +143,7 @@ E20.py:90:18: E202 [*] Whitespace before ']'
 90    |-f"{[ { {'a': 1} } ]}"
    90 |+f"{[ { {'a': 1} }]}"
 91 91 | f"normal { {f"{ { [1, 2] } }" } } normal"
+92 92 | 
+93 93 | #: Okay
 
 
diff --git a/.../src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E203_E20.py.snap b/.../src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E203_E20.py.snap
@@ -126,4 +126,99 @@ E20.py:71:13: E203 [*] Whitespace before ','
 73 73 | if x == 4:
 74 74 |     print x, y
 
+E20.py:100:19: E203 [*] Whitespace before ':'
+    |
+ 99 | #: E203:1:19
+100 | ham{lower + offset : upper + offset}
+    |                   ^ E203
+101 | 
+102 | #: E203:1:19
+    |
+    = help: Remove whitespace before ':'
+
+ℹ Safe fix
+97  97  | ham[(lower + offset) : upper + offset]
+98  98  | 
+99  99  | #: E203:1:19
+100     |-ham{lower + offset : upper + offset}
+    100 |+ham{lower + offset: upper + offset}
+101 101 | 
+102 102 | #: E203:1:19
+103 103 | ham[lower + offset  : upper + offset]
+
+E20.py:103:19: E203 [*] Whitespace before ':'
+    |
+102 | #: E203:1:19
+103 | ham[lower + offset  : upper + offset]
+    |                   ^^ E203
+104 | 
+105 | #: Okay
+    |
+    = help: Remove whitespace before ':'
+
+ℹ Safe fix
+100 100 | ham{lower + offset : upper + offset}
+101 101 | 
+102 102 | #: E203:1:19
+103     |-ham[lower + offset  : upper + offset]
+    103 |+ham[lower + offset: upper + offset]
+104 104 | 
+105 105 | #: Okay
+106 106 | release_lines = history_file_lines[history_file_lines.index('## Unreleased') + 1: -1]
+
+E20.py:131:21: E203 [*] Whitespace before ':'
+    |
+130 | #: E203:1:21
+131 | ham[lower + offset : : upper + offset]
+    |                     ^ E203
+132 | 
+133 | #: E203:1:20
+    |
+    = help: Remove whitespace before ':'
+
+ℹ Safe fix
+128 128 | ham[lower + offset::upper + offset]
+129 129 | 
+130 130 | #: E203:1:21
+131     |-ham[lower + offset : : upper + offset]
+    131 |+ham[lower + offset :: upper + offset]
+132 132 | 
+133 133 | #: E203:1:20
+134 134 | ham[lower + offset: :upper + offset]
+
+E20.py:134:20: E203 [*] Whitespace before ':'
+    |
+133 | #: E203:1:20
+134 | ham[lower + offset: :upper + offset]
+    |                    ^ E203
+135 | 
+136 | #: E203:1:20
+    |
+    = help: Remove whitespace before ':'
+
+ℹ Safe fix
+131 131 | ham[lower + offset : : upper + offset]
+132 132 | 
+133 133 | #: E203:1:20
+134     |-ham[lower + offset: :upper + offset]
+    134 |+ham[lower + offset::upper + offset]
+135 135 | 
+136 136 | #: E203:1:20
+137 137 | ham[{lower + offset : upper + offset} : upper + offset]
+
+E20.py:137:20: E203 [*] Whitespace before ':'
+    |
+136 | #: E203:1:20
+137 | ham[{lower + offset : upper + offset} : upper + offset]
+    |                    ^ E203
+    |
+    = help: Remove whitespace before ':'
+
+ℹ Safe fix
+134 134 | ham[lower + offset: :upper + offset]
+135 135 | 
+136 136 | #: E203:1:20
+137     |-ham[{lower + offset : upper + offset} : upper + offset]
+    137 |+ham[{lower + offset: upper + offset} : upper + offset]
+