Adapted fix to work identical to format (#10999)

## Summary

The fix for E203 now produces the same result as ruff format in cases
where a slice ends on a colon and the closing square bracket is on the
following line.

Refers to https://github.com/astral-sh/ruff/issues/10973

## Test Plan

The minimal reproduction case in the ticket was added as test case
producing no error. Additional cases with multiple spaces or a tab
before the colon where added to make sure that the rule still finds
these.

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

@@ -75,6 +75,32 @@ if x == 4:
     x, y = y, x
 a[b1, :] == a[b1, ...]
 b = a[:, b1]
+
+#: E203 linebreak before ]
+predictions = predictions[
+    len(past_covariates) // datamodule.hparams["downsample"] :
+]
+
+#: E203 multi whitespace before :
+predictions = predictions[
+    len(past_covariates) // datamodule.hparams["downsample"]  :
+]
+
+#: E203 tab before :
+predictions = predictions[
+    len(past_covariates) // datamodule.hparams["downsample"]	:
+]
+
+#: E203 single whitespace before : with line a comment
+predictions = predictions[
+    len(past_covariates) // datamodule.hparams["downsample"] :  # Just some comment
+]
+
+#: E203 multi whitespace before : with line a comment
+predictions = predictions[
+    len(past_covariates) // datamodule.hparams["downsample"]  :  # Just some comment
+]
+
 #:
 
 #: E201:1:6

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

@@ -229,6 +229,32 @@ pub(crate) fn extraneous_whitespace(line: &LogicalLine, context: &mut LogicalLin
                                         )));
                                         context.push_diagnostic(diagnostic);
                                     }
+                                } else if iter.peek().is_some_and(|token| {
+                                    matches!(
+                                        token.kind(),
+                                        TokenKind::NonLogicalNewline | TokenKind::Comment
+                                    )
+                                }) {
+                                    // Allow [
+                                    //      long_expression_calculating_the_index() :
+                                    // ]
+                                    // But not [
+                                    //      long_expression_calculating_the_index()  :
+                                    // ]
+                                    // distinct from the above case, because ruff format produces a
+                                    // whitespace before the colon and so should the fix
+                                    if let (Whitespace::Many | Whitespace::Tab, offset) = whitespace
+                                    {
+                                        let mut diagnostic = Diagnostic::new(
+                                            WhitespaceBeforePunctuation { symbol },
+                                            TextRange::at(token.start() - offset, offset),
+                                        );
+                                        diagnostic.set_fix(Fix::safe_edits(
+                                            Edit::range_deletion(diagnostic.range()),
+                                            [Edit::insertion(" ".into(), token.start() - offset)],
+                                        ));
+                                        context.push_diagnostic(diagnostic);
+                                    }
                                 } else {
                                     // Allow, e.g., `foo[1:2]` or `foo[1 : 2]` or `foo[1 :: 2]`.
                                     let token = iter

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

@@ -124,64 +124,62 @@ E20.py:12:15: E201 [*] Whitespace after '{'
 14 14 | spam(ham[1], {eggs: 2})
 15 15 | #:
 
-E20.py:81:6: E201 [*] Whitespace after '['
+E20.py:107:6: E201 [*] Whitespace after '['
     |
-80 | #: E201:1:6
+106 | #: E201:1:6
-81 | spam[ ~ham]
+107 | spam[ ~ham]
     |      ^ E201
-82 | 
+108 | 
-83 | #: Okay
+109 | #: Okay
     |
     = help: Remove whitespace before '['
 
 ℹ Safe fix
-78 78 | #:
+104 104 | #:
-79 79 | 
+105 105 | 
-80 80 | #: E201:1:6
+106 106 | #: E201:1:6
-81    |-spam[ ~ham]
+107     |-spam[ ~ham]
-   81 |+spam[~ham]
+    107 |+spam[~ham]
-82 82 | 
+108 108 | 
-83 83 | #: Okay
+109 109 | #: Okay
-84 84 | x = [  #
+110 110 | x = [  #
 
-E20.py:90:5: E201 [*] Whitespace after '['
+E20.py:116:5: E201 [*] Whitespace after '['
     |
-88 | # F-strings
+114 | # F-strings
-89 | f"{ {'a': 1} }"
+115 | f"{ {'a': 1} }"
-90 | f"{[ { {'a': 1} } ]}"
+116 | f"{[ { {'a': 1} } ]}"
     |     ^ E201
-91 | f"normal { {f"{ { [1, 2] } }" } } normal"
+117 | f"normal { {f"{ { [1, 2] } }" } } normal"
     |
     = help: Remove whitespace before '['
 
 ℹ Safe fix
-87 87 | 
+113 113 | 
-88 88 | # F-strings
+114 114 | # F-strings
-89 89 | f"{ {'a': 1} }"
+115 115 | f"{ {'a': 1} }"
-90    |-f"{[ { {'a': 1} } ]}"
+116     |-f"{[ { {'a': 1} } ]}"
-   90 |+f"{[{ {'a': 1} } ]}"
+    116 |+f"{[{ {'a': 1} } ]}"
-91 91 | f"normal { {f"{ { [1, 2] } }" } } normal"
+117 117 | f"normal { {f"{ { [1, 2] } }" } } normal"
-92 92 | 
+118 118 | 
-93 93 | #: Okay
+119 119 | #: Okay
 
-E20.py:119:5: E201 [*] Whitespace after '['
+E20.py:145:5: E201 [*] Whitespace after '['
     |
-118 | #: E201:1:5
+144 | #: E201:1:5
-119 | ham[ : upper]
+145 | ham[ : upper]
     |     ^ E201
-120 | 
+146 | 
-121 | #: Okay
+147 | #: Okay
     |
     = help: Remove whitespace before '['
 
 ℹ Safe fix
-116 116 | ham[lower + offset : upper + offset]
+142 142 | ham[lower + offset : upper + offset]
-117 117 | 
+143 143 | 
-118 118 | #: E201:1:5
+144 144 | #: E201:1:5
-119     |-ham[ : upper]
+145     |-ham[ : upper]
-    119 |+ham[: upper]
+    145 |+ham[: upper]
-120 120 | 
+146 146 | 
-121 121 | #: Okay
+147 147 | #: Okay
-122 122 | ham[lower + offset :: upper + offset]
+148 148 | ham[lower + offset :: upper + offset]
-
-

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

@@ -126,44 +126,42 @@ E20.py:29:11: E202 [*] Whitespace before ']'
 31 31 | spam(ham[1], {eggs: 2})
 32 32 | 
 
-E20.py:90:18: E202 [*] Whitespace before ']'
+E20.py:116:18: E202 [*] Whitespace before ']'
     |
-88 | # F-strings
+114 | # F-strings
-89 | f"{ {'a': 1} }"
+115 | f"{ {'a': 1} }"
-90 | f"{[ { {'a': 1} } ]}"
+116 | f"{[ { {'a': 1} } ]}"
     |                  ^ E202
-91 | f"normal { {f"{ { [1, 2] } }" } } normal"
+117 | f"normal { {f"{ { [1, 2] } }" } } normal"
     |
     = help: Remove whitespace before ']'
 
 ℹ Safe fix
-87 87 | 
+113 113 | 
-88 88 | # F-strings
+114 114 | # F-strings
-89 89 | f"{ {'a': 1} }"
+115 115 | f"{ {'a': 1} }"
-90    |-f"{[ { {'a': 1} } ]}"
+116     |-f"{[ { {'a': 1} } ]}"
-   90 |+f"{[ { {'a': 1} }]}"
+    116 |+f"{[ { {'a': 1} }]}"
-91 91 | f"normal { {f"{ { [1, 2] } }" } } normal"
+117 117 | f"normal { {f"{ { [1, 2] } }" } } normal"
-92 92 | 
+118 118 | 
-93 93 | #: Okay
+119 119 | #: Okay
 
-E20.py:146:12: E202 [*] Whitespace before ']'
+E20.py:172:12: E202 [*] Whitespace before ']'
     |
-145 | #: E202:1:12
+171 | #: E202:1:12
-146 | ham[upper : ]
+172 | ham[upper : ]
     |            ^ E202
-147 | 
+173 | 
-148 | #: E203:1:10
+174 | #: E203:1:10
     |
     = help: Remove whitespace before ']'
 
 ℹ Safe fix
-143 143 | ham[upper :]
+169 169 | ham[upper :]
-144 144 | 
+170 170 | 
-145 145 | #: E202:1:12
+171 171 | #: E202:1:12
-146     |-ham[upper : ]
+172     |-ham[upper : ]
-    146 |+ham[upper :]
+    172 |+ham[upper :]
-147 147 | 
+173 173 | 
-148 148 | #: E203:1:10
+174 174 | #: E203:1:10
-149 149 | ham[upper  :]
+175 175 | ham[upper  :]
-
-

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

@@ -126,157 +126,217 @@ E20.py:71:13: E203 [*] Whitespace before ','
 73 73 | if x == 4:
 74 74 |     print x, y
 
-E20.py:100:19: E203 [*] Whitespace before ':'
+E20.py:86:61: E203 [*] Whitespace before ':'
    |
- 99 | #: E203:1:19
+84 | #: E203 multi whitespace before :
-100 | ham{lower + offset : upper + offset}
+85 | predictions = predictions[
-    |                   ^ E203
+86 |     len(past_covariates) // datamodule.hparams["downsample"]  :
-101 | 
+   |                                                             ^^ E203
-102 | #: E203:1:19
+87 | ]
+   |
+   = help: Remove whitespace before ':'
+
+ℹ Safe fix
+83 83 | 
+84 84 | #: E203 multi whitespace before :
+85 85 | predictions = predictions[
+86    |-    len(past_covariates) // datamodule.hparams["downsample"]  :
+   86 |+    len(past_covariates) // datamodule.hparams["downsample"] :
+87 87 | ]
+88 88 | 
+89 89 | #: E203 tab before :
+
+E20.py:91:61: E203 [*] Whitespace before ':'
+   |
+89 | #: E203 tab before :
+90 | predictions = predictions[
+91 |     len(past_covariates) // datamodule.hparams["downsample"]    :
+   |                                                             ^^^^ E203
+92 | ]
+   |
+   = help: Remove whitespace before ':'
+
+ℹ Safe fix
+88 88 | 
+89 89 | #: E203 tab before :
+90 90 | predictions = predictions[
+91    |-    len(past_covariates) // datamodule.hparams["downsample"]	:
+   91 |+    len(past_covariates) // datamodule.hparams["downsample"] :
+92 92 | ]
+93 93 | 
+94 94 | #: E203 single whitespace before : with line a comment
+
+E20.py:101:61: E203 [*] Whitespace before ':'
+    |
+ 99 | #: E203 multi whitespace before : with line a comment
+100 | predictions = predictions[
+101 |     len(past_covariates) // datamodule.hparams["downsample"]  :  # Just some comment
+    |                                                             ^^ E203
+102 | ]
     |
     = help: Remove whitespace before ':'
 
 ℹ Safe fix
-97  97  | ham[(lower + offset) : upper + offset]
 98  98  | 
-99  99  | #: E203:1:19
+99  99  | #: E203 multi whitespace before : with line a comment
-100     |-ham{lower + offset : upper + offset}
+100 100 | predictions = predictions[
-    100 |+ham{lower + offset: upper + offset}
+101     |-    len(past_covariates) // datamodule.hparams["downsample"]  :  # Just some comment
-101 101 | 
+    101 |+    len(past_covariates) // datamodule.hparams["downsample"] :  # Just some comment
-102 102 | #: E203:1:19
+102 102 | ]
-103 103 | ham[lower + offset  : upper + offset]
+103 103 | 
+104 104 | #:
 
-E20.py:103:19: E203 [*] Whitespace before ':'
+E20.py:126:19: E203 [*] Whitespace before ':'
     |
-102 | #: E203:1:19
+125 | #: E203:1:19
-103 | ham[lower + offset  : upper + offset]
+126 | 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 | 
+127 | 
-133 | #: E203:1:20
+128 | #: E203:1:19
     |
     = help: Remove whitespace before ':'
 
 ℹ Safe fix
-128 128 | ham[lower + offset::upper + offset]
+123 123 | ham[(lower + offset) : upper + offset]
-129 129 | 
+124 124 | 
-130 130 | #: E203:1:21
+125 125 | #: E203:1:19
-131     |-ham[lower + offset : : upper + offset]
+126     |-ham{lower + offset : upper + offset}
-    131 |+ham[lower + offset :: upper + offset]
+    126 |+ham{lower + offset: upper + offset}
-132 132 | 
+127 127 | 
-133 133 | #: E203:1:20
+128 128 | #: E203:1:19
-134 134 | ham[lower + offset: :upper + offset]
+129 129 | ham[lower + offset  : upper + offset]
 
-E20.py:134:20: E203 [*] Whitespace before ':'
+E20.py:129:19: E203 [*] Whitespace before ':'
     |
-133 | #: E203:1:20
+128 | #: E203:1:19
-134 | ham[lower + offset: :upper + offset]
+129 | ham[lower + offset  : upper + offset]
+    |                   ^^ E203
+130 | 
+131 | #: Okay
+    |
+    = help: Remove whitespace before ':'
+
+ℹ Safe fix
+126 126 | ham{lower + offset : upper + offset}
+127 127 | 
+128 128 | #: E203:1:19
+129     |-ham[lower + offset  : upper + offset]
+    129 |+ham[lower + offset: upper + offset]
+130 130 | 
+131 131 | #: Okay
+132 132 | release_lines = history_file_lines[history_file_lines.index('## Unreleased') + 1: -1]
+
+E20.py:157:21: E203 [*] Whitespace before ':'
+    |
+156 | #: E203:1:21
+157 | ham[lower + offset : : upper + offset]
     |                     ^ E203
-135 | 
+158 | 
-136 | #: E203:1:20
+159 | #: E203:1:20
     |
     = help: Remove whitespace before ':'
 
 ℹ Safe fix
-131 131 | ham[lower + offset : : upper + offset]
+154 154 | ham[lower + offset::upper + offset]
-132 132 | 
+155 155 | 
-133 133 | #: E203:1:20
+156 156 | #: E203:1:21
-134     |-ham[lower + offset: :upper + offset]
+157     |-ham[lower + offset : : upper + offset]
-    134 |+ham[lower + offset::upper + offset]
+    157 |+ham[lower + offset :: upper + offset]
-135 135 | 
+158 158 | 
-136 136 | #: E203:1:20
+159 159 | #: E203:1:20
-137 137 | ham[{lower + offset : upper + offset} : upper + offset]
+160 160 | ham[lower + offset: :upper + offset]
 
-E20.py:137:20: E203 [*] Whitespace before ':'
+E20.py:160:20: E203 [*] Whitespace before ':'
     |
-136 | #: E203:1:20
+159 | #: E203:1:20
-137 | ham[{lower + offset : upper + offset} : upper + offset]
+160 | ham[lower + offset: :upper + offset]
     |                    ^ E203
-138 | 
+161 | 
-139 | #: Okay
+162 | #: E203:1:20
     |
     = help: Remove whitespace before ':'
 
 ℹ Safe fix
-134 134 | ham[lower + offset: :upper + offset]
+157 157 | ham[lower + offset : : upper + offset]
-135 135 | 
+158 158 | 
-136 136 | #: E203:1:20
+159 159 | #: E203:1:20
-137     |-ham[{lower + offset : upper + offset} : upper + offset]
+160     |-ham[lower + offset: :upper + offset]
-    137 |+ham[{lower + offset: upper + offset} : upper + offset]
+    160 |+ham[lower + offset::upper + offset]
-138 138 | 
+161 161 | 
-139 139 | #: Okay
+162 162 | #: E203:1:20
-140 140 | ham[upper:]
+163 163 | ham[{lower + offset : upper + offset} : upper + offset]
 
-E20.py:149:10: E203 [*] Whitespace before ':'
+E20.py:163:20: E203 [*] Whitespace before ':'
     |
-148 | #: E203:1:10
+162 | #: E203:1:20
-149 | ham[upper  :]
+163 | ham[{lower + offset : upper + offset} : upper + offset]
+    |                    ^ E203
+164 | 
+165 | #: Okay
+    |
+    = help: Remove whitespace before ':'
+
+ℹ Safe fix
+160 160 | ham[lower + offset: :upper + offset]
+161 161 | 
+162 162 | #: E203:1:20
+163     |-ham[{lower + offset : upper + offset} : upper + offset]
+    163 |+ham[{lower + offset: upper + offset} : upper + offset]
+164 164 | 
+165 165 | #: Okay
+166 166 | ham[upper:]
+
+E20.py:175:10: E203 [*] Whitespace before ':'
+    |
+174 | #: E203:1:10
+175 | ham[upper  :]
     |          ^^ E203
-150 | 
+176 | 
-151 | #: Okay
+177 | #: Okay
     |
     = help: Remove whitespace before ':'
 
 ℹ Safe fix
-146 146 | ham[upper : ]
+172 172 | ham[upper : ]
-147 147 | 
+173 173 | 
-148 148 | #: E203:1:10
+174 174 | #: E203:1:10
-149     |-ham[upper  :]
+175     |-ham[upper  :]
-    149 |+ham[upper:]
+    175 |+ham[upper:]
-150 150 | 
+176 176 | 
-151 151 | #: Okay
+177 177 | #: Okay
-152 152 | ham[lower +1 :, "columnname"]
+178 178 | ham[lower +1 :, "columnname"]
 
-E20.py:155:14: E203 [*] Whitespace before ':'
+E20.py:181:14: E203 [*] Whitespace before ':'
     |
-154 | #: E203:1:13
+180 | #: E203:1:13
-155 | ham[lower + 1  :, "columnname"]
+181 | ham[lower + 1  :, "columnname"]
     |              ^^ E203
-156 | 
+182 | 
-157 | #: Okay
+183 | #: Okay
     |
     = help: Remove whitespace before ':'
 
 ℹ Safe fix
-152 152 | ham[lower +1 :, "columnname"]
+178 178 | ham[lower +1 :, "columnname"]
-153 153 | 
+179 179 | 
-154 154 | #: E203:1:13
+180 180 | #: E203:1:13
-155     |-ham[lower + 1  :, "columnname"]
+181     |-ham[lower + 1  :, "columnname"]
-    155 |+ham[lower + 1:, "columnname"]
+    181 |+ham[lower + 1:, "columnname"]
-156 156 | 
+182 182 | 
-157 157 | #: Okay
+183 183 | #: Okay
-158 158 | f"{ham[lower +1 :, "columnname"]}"
+184 184 | f"{ham[lower +1 :, "columnname"]}"
 
-E20.py:161:17: E203 [*] Whitespace before ':'
+E20.py:187:17: E203 [*] Whitespace before ':'
     |
-160 | #: E203:1:13
+186 | #: E203:1:13
-161 | f"{ham[lower + 1  :, "columnname"]}"
+187 | f"{ham[lower + 1  :, "columnname"]}"
     |                 ^^ E203
     |
     = help: Remove whitespace before ':'
 
 ℹ Safe fix
-158 158 | f"{ham[lower +1 :, "columnname"]}"
+184 184 | f"{ham[lower +1 :, "columnname"]}"
-159 159 | 
+185 185 | 
-160 160 | #: E203:1:13
+186 186 | #: E203:1:13
-161     |-f"{ham[lower + 1  :, "columnname"]}"
+187     |-f"{ham[lower + 1  :, "columnname"]}"
-    161 |+f"{ham[lower + 1:, "columnname"]}"
+    187 |+f"{ham[lower + 1:, "columnname"]}"