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 #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 @@
     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 '['
-   |
-80 | #: E201:1:6
-81 | spam[ ~ham]
-   |      ^ E201
-82 | 
-83 | #: Okay
-   |
-   = help: Remove whitespace before '['
+E20.py:107:6: E201 [*] Whitespace after '['
+    |
+106 | #: E201:1:6
+107 | spam[ ~ham]
+    |      ^ E201
+108 | 
+109 | #: Okay
+    |
+    = help: Remove whitespace before '['
 
 ℹ Safe fix
-78 78 | #:
-79 79 | 
-80 80 | #: E201:1:6
-81    |-spam[ ~ham]
-   81 |+spam[~ham]
-82 82 | 
-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 '['
+104 104 | #:
+105 105 | 
+106 106 | #: E201:1:6
+107     |-spam[ ~ham]
+    107 |+spam[~ham]
+108 108 | 
+109 109 | #: Okay
+110 110 | x = [  #
+
+E20.py:116:5: E201 [*] Whitespace after '['
+    |
+114 | # F-strings
+115 | f"{ {'a': 1} }"
+116 | f"{[ { {'a': 1} } ]}"
+    |     ^ E201
+117 | f"normal { {f"{ { [1, 2] } }" } } normal"
+    |
+    = help: Remove whitespace before '['
 
 ℹ Safe 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"
-92 92 | 
-93 93 | #: Okay
-
-E20.py:119:5: E201 [*] Whitespace after '['
+113 113 | 
+114 114 | # F-strings
+115 115 | f"{ {'a': 1} }"
+116     |-f"{[ { {'a': 1} } ]}"
+    116 |+f"{[{ {'a': 1} } ]}"
+117 117 | f"normal { {f"{ { [1, 2] } }" } } normal"
+118 118 | 
+119 119 | #: Okay
+
+E20.py:145:5: E201 [*] Whitespace after '['
     |
-118 | #: E201:1:5
-119 | ham[ : upper]
+144 | #: E201:1:5
+145 | ham[ : upper]
     |     ^ E201
-120 | 
-121 | #: Okay
+146 | 
+147 | #: 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]
-
-
+142 142 | ham[lower + offset : upper + offset]
+143 143 | 
+144 144 | #: E201:1:5
+145     |-ham[ : upper]
+    145 |+ham[: upper]
+146 146 | 
+147 147 | #: Okay
+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 ']'
-   |
-88 | # F-strings
-89 | f"{ {'a': 1} }"
-90 | f"{[ { {'a': 1} } ]}"
-   |                  ^ E202
-91 | f"normal { {f"{ { [1, 2] } }" } } normal"
-   |
-   = help: Remove whitespace before ']'
+E20.py:116:18: E202 [*] Whitespace before ']'
+    |
+114 | # F-strings
+115 | f"{ {'a': 1} }"
+116 | f"{[ { {'a': 1} } ]}"
+    |                  ^ E202
+117 | f"normal { {f"{ { [1, 2] } }" } } normal"
+    |
+    = help: Remove whitespace before ']'
 
 ℹ Safe 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"
-92 92 | 
-93 93 | #: Okay
-
-E20.py:146:12: E202 [*] Whitespace before ']'
+113 113 | 
+114 114 | # F-strings
+115 115 | f"{ {'a': 1} }"
+116     |-f"{[ { {'a': 1} } ]}"
+    116 |+f"{[ { {'a': 1} }]}"
+117 117 | f"normal { {f"{ { [1, 2] } }" } } normal"
+118 118 | 
+119 119 | #: Okay
+
+E20.py:172:12: E202 [*] Whitespace before ']'
     |
-145 | #: E202:1:12
-146 | ham[upper : ]
+171 | #: E202:1:12
+172 | ham[upper : ]
     |            ^ E202
-147 | 
-148 | #: E203:1:10
+173 | 
+174 | #: E203:1:10
     |
     = help: Remove whitespace before ']'
 
 ℹ Safe fix
-143 143 | ham[upper :]
-144 144 | 
-145 145 | #: E202:1:12
-146     |-ham[upper : ]
-    146 |+ham[upper :]
-147 147 | 
-148 148 | #: E203:1:10
-149 149 | ham[upper  :]
-
-
+169 169 | ham[upper :]
+170 170 | 
+171 171 | #: E202:1:12
+172     |-ham[upper : ]
+    172 |+ham[upper :]
+173 173 | 
+174 174 | #: E203:1:10
+175 175 | ham[upper  :]