Add a fix

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

@@ -24,4 +24,11 @@ def bar(self):
     # E204
     @ foo
     def baz(self):
-        print('baz')
+        print('baz')
+
+
+# E204
+@ \
+foo
+def baz():
+    print('baz')

crates/ruff_linter/src/rules/pycodestyle/mod.rs

@@ -58,6 +58,7 @@ mod tests {
     #[test_case(Rule::TypeComparison, Path::new("E721.py"))]
     #[test_case(Rule::UselessSemicolon, Path::new("E70.py"))]
     #[test_case(Rule::UselessSemicolon, Path::new("E703.ipynb"))]
+    #[test_case(Rule::WhitespaceAfterDecorator, Path::new("E204.py"))]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());
         let diagnostics = test_path(

crates/ruff_linter/src/rules/pycodestyle/rules/whitespace_after_decorator.rs

@@ -1,30 +1,34 @@
-use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic};
+use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::Decorator;
 use ruff_python_trivia::is_python_whitespace;
-use ruff_text_size::{TextLen, TextRange};
+use ruff_text_size::{Ranged, TextRange, TextSize};
 
 use crate::checkers::ast::Checker;
 
 /// ## What it does
-/// Checks for whitespace after a decorator.
+/// Checks for trailing whitespace after a decorator's opening `@`.
 ///
 /// ## Why is this bad?
-/// Whitespace after a decorator is not PEP8 compliant.
+/// Including whitespace after the `@` symbol is not compliant with
+/// [PEP 8].
 ///
 /// ## Example
 ///
 /// ```python
 /// @ decorator
-/// def foo():
+/// def func():
 ///    pass
 /// ```
+///
 /// Use instead:
-/// ``` python
+/// ```python
 /// @decorator
-/// def foo():
+/// def func():
 ///   pass
 /// ```
+///
+/// [PEP 8]: https://peps.python.org/pep-0008/#maximum-line-length
 
 #[violation]
 pub struct WhitespaceAfterDecorator;
@@ -36,29 +40,32 @@ impl AlwaysFixableViolation for WhitespaceAfterDecorator {
     }
 
     fn fix_title(&self) -> String {
-        "Remove whitespace after decorator".to_string()
+        "Remove whitespace".to_string()
     }
 }
 
 /// E204
 pub(crate) fn whitespace_after_decorator(checker: &mut Checker, decorator_list: &[Decorator]) {
-    // Get the locator from the checker
-    let locator = checker.locator();
-
-    // Iterate over the list of decorators
     for decorator in decorator_list {
-        // Obtain the text of the decorator using lactor.slice(decorator)
-        let decorator_text = locator.slice(decorator);
+        let decorator_text = checker.locator().slice(decorator);
+
+        // Determine whether the `@` is followed by whitespace.
+        if let Some(trailing) = decorator_text.strip_prefix('@') {
+            // Collect the whitespace characters after the `@`.
+            if trailing.chars().next().is_some_and(is_python_whitespace) {
+                let end = trailing
+                    .chars()
+                    .position(|c| !(is_python_whitespace(c) || matches!(c, '\n' | '\r' | '\\')))
+                    .unwrap_or(trailing.len());
 
-        // Get the text after the @ symbol
-        let after_at = &decorator_text[1..];
+                let start = decorator.start() + TextSize::from(1);
+                let end = start + TextSize::from(end as u32);
+                let range = TextRange::new(start, end);
 
-        // Check if there is a whitespace after the @ symbol by using is_python_whitespace
-        if is_python_whitespace(after_at.chars().next().unwrap()) {
-            let range = TextRange::empty(locator.contents().text_len());
-            checker
-                .diagnostics
-                .push(Diagnostic::new(WhitespaceAfterDecorator, range));
+                let mut diagnostic = Diagnostic::new(WhitespaceAfterDecorator, range);
+                diagnostic.set_fix(Fix::safe_edit(Edit::range_deletion(range)));
+                checker.diagnostics.push(diagnostic);
+            }
         }
     }
 }

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

@@ -0,0 +1,63 @@
+---
+source: crates/ruff_linter/src/rules/pycodestyle/mod.rs
+---
+E204.py:14:1: E204 [*] Whitespace after decorator
+   |
+13 | # E204
+14 | @ foo
+   | ^^^^^ E204
+15 | def baz():
+16 |     print('baz')
+   |
+   = help: Remove whitespace
+
+ℹ Safe fix
+11 11 |     print('bar')
+12 12 | 
+13 13 | # E204
+14    |-@ foo
+   14 |+@foo
+15 15 | def baz():
+16 16 |     print('baz')
+17 17 | 
+
+E204.py:25:5: E204 [*] Whitespace after decorator
+   |
+24 |     # E204
+25 |     @ foo
+   |     ^^^^^ E204
+26 |     def baz(self):
+27 |         print('baz')
+   |
+   = help: Remove whitespace
+
+ℹ Safe fix
+22 22 |         print('bar')
+23 23 | 
+24 24 |     # E204
+25    |-    @ foo
+   25 |+    @foo
+26 26 |     def baz(self):
+27 27 |         print('baz')
+28 28 | 
+
+E204.py:31:1: E204 [*] Whitespace after decorator
+   |
+30 |   # E204
+31 | / @ \
+32 | | foo
+   | |___^ E204
+33 |   def baz():
+34 |       print('baz')
+   |
+   = help: Remove whitespace
+
+ℹ Safe fix
+28 28 | 
+29 29 | 
+30 30 | # E204
+31    |-@ \
+32    |-foo
+   31 |+@foo
+33 32 | def baz():
+34 33 |     print('baz')

crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E204_E204.py.snap.new

@@ -0,0 +1,65 @@
+---
+source: crates/ruff_linter/src/rules/pycodestyle/mod.rs
+assertion_line: 68
+---
+E204.py:14:2: E204 [*] Whitespace after decorator
+   |
+13 | # E204
+14 | @ foo
+   |  ^ E204
+15 | def baz():
+16 |     print('baz')
+   |
+   = help: Remove whitespace
+
+ℹ Safe fix
+11 11 |     print('bar')
+12 12 | 
+13 13 | # E204
+14    |-@ foo
+   14 |+@foo
+15 15 | def baz():
+16 16 |     print('baz')
+17 17 | 
+
+E204.py:25:6: E204 [*] Whitespace after decorator
+   |
+24 |     # E204
+25 |     @ foo
+   |      ^ E204
+26 |     def baz(self):
+27 |         print('baz')
+   |
+   = help: Remove whitespace
+
+ℹ Safe fix
+22 22 |         print('bar')
+23 23 | 
+24 24 |     # E204
+25    |-    @ foo
+   25 |+    @foo
+26 26 |     def baz(self):
+27 27 |         print('baz')
+28 28 | 
+
+E204.py:31:2: E204 [*] Whitespace after decorator
+   |
+30 |   # E204
+31 |   @ \
+   |  __^
+32 | | foo
+   | |_^ E204
+33 |   def baz():
+34 |       print('baz')
+   |
+   = help: Remove whitespace
+
+ℹ Safe fix
+28 28 | 
+29 29 | 
+30 30 | # E204
+31    |-@ \
+32    |-foo
+   31 |+@foo
+33 32 | def baz():
+34 33 |     print('baz')