[pycodestyle] Whitespace after decorator (E204) (#12140)

## Summary

<!-- What's the purpose of the change? What does it do, and why? -->
This is the implementation for the new rule of `pycodestyle (E204)`. It
follows the guidlines described in the contributing site, and as such it
has a new file named `whitespace_after_decorator.rs`, a new test file
called `E204.py`, and as such invokes the `function` in the `AST
statement checker` for functions and functions in classes. Linking #2402
because it has all the pycodestyle rules.

## Test Plan

<!-- How was it tested? -->
The file E204.py, has a `decorator` defined called wrapper, and this
decorator is used for 2 cases. The first one is when a `function` which
has a `decorator` is called in the file, and the second one is when
there is a `class` and 2 `methods` are defined for the `class` with a
`decorator` attached it.

Test file:

``` python
def foo(fun):
    def wrapper():
        print('before')
        fun()
        print('after')
    return wrapper

# No error
@foo
def bar():
    print('bar')

# E204
@ foo
def baz():
    print('baz')

class Test:
    # No error
    @foo
    def bar(self):
        print('bar')

    # E204
    @ foo
    def baz(self):
        print('baz')
```

I am still new to rust and any suggestion is appreciated. Specially with
the way im using native ruff utilities.

---------

Co-authored-by: Charlie Marsh <[email protected]>

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

@@ -0,0 +1,34 @@
+def foo(fun):
+    def wrapper():
+        print('before')
+        fun()
+        print('after')
+    return wrapper
+
+# No error
+@foo
+def bar():
+    print('bar')
+
+# E204
+@ foo
+def baz():
+    print('baz')
+
+class Test:
+    # No error
+    @foo
+    def bar(self):
+        print('bar')
+
+    # E204
+    @ foo
+    def baz(self):
+        print('baz')
+
+
+# E204
+@ \
+foo
+def baz():
+    print('baz')

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -368,6 +368,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.enabled(Rule::UnusedAsync) {
                 ruff::rules::unused_async(checker, function_def);
             }
+            if checker.enabled(Rule::WhitespaceAfterDecorator) {
+                pycodestyle::rules::whitespace_after_decorator(checker, decorator_list);
+            }
         }
         Stmt::Return(_) => {
             if checker.enabled(Rule::ReturnOutsideFunction) {
@@ -531,6 +534,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.enabled(Rule::MetaClassABCMeta) {
                 refurb::rules::metaclass_abcmeta(checker, class_def);
             }
+            if checker.enabled(Rule::WhitespaceAfterDecorator) {
+                pycodestyle::rules::whitespace_after_decorator(checker, decorator_list);
+            }
         }
         Stmt::Import(ast::StmtImport { names, range: _ }) => {
             if checker.enabled(Rule::MultipleImportsOnOneLine) {

crates/ruff_linter/src/codes.rs

@@ -78,6 +78,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Pycodestyle, "E201") => (RuleGroup::Preview, rules::pycodestyle::rules::logical_lines::WhitespaceAfterOpenBracket),
         (Pycodestyle, "E202") => (RuleGroup::Preview, rules::pycodestyle::rules::logical_lines::WhitespaceBeforeCloseBracket),
         (Pycodestyle, "E203") => (RuleGroup::Preview, rules::pycodestyle::rules::logical_lines::WhitespaceBeforePunctuation),
+        (Pycodestyle, "E204") => (RuleGroup::Preview, rules::pycodestyle::rules::WhitespaceAfterDecorator),
         (Pycodestyle, "E211") => (RuleGroup::Preview, rules::pycodestyle::rules::logical_lines::WhitespaceBeforeParameters),
         (Pycodestyle, "E221") => (RuleGroup::Preview, rules::pycodestyle::rules::logical_lines::MultipleSpacesBeforeOperator),
         (Pycodestyle, "E222") => (RuleGroup::Preview, rules::pycodestyle::rules::logical_lines::MultipleSpacesAfterOperator),

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/mod.rs

@@ -20,6 +20,7 @@ pub(crate) use tab_indentation::*;
 pub(crate) use too_many_newlines_at_end_of_file::*;
 pub(crate) use trailing_whitespace::*;
 pub(crate) use type_comparison::*;
+pub(crate) use whitespace_after_decorator::*;
 
 mod ambiguous_class_name;
 mod ambiguous_function_name;
@@ -43,3 +44,4 @@ mod tab_indentation;
 mod too_many_newlines_at_end_of_file;
 mod trailing_whitespace;
 mod type_comparison;
+mod whitespace_after_decorator;

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

@@ -0,0 +1,71 @@
+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::{Ranged, TextRange, TextSize};
+
+use crate::checkers::ast::Checker;
+
+/// ## What it does
+/// Checks for trailing whitespace after a decorator's opening `@`.
+///
+/// ## Why is this bad?
+/// Including whitespace after the `@` symbol is not compliant with
+/// [PEP 8].
+///
+/// ## Example
+///
+/// ```python
+/// @ decorator
+/// def func():
+///    pass
+/// ```
+///
+/// Use instead:
+/// ```python
+/// @decorator
+/// def func():
+///   pass
+/// ```
+///
+/// [PEP 8]: https://peps.python.org/pep-0008/#maximum-line-length
+
+#[violation]
+pub struct WhitespaceAfterDecorator;
+
+impl AlwaysFixableViolation for WhitespaceAfterDecorator {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        format!("Whitespace after decorator")
+    }
+
+    fn fix_title(&self) -> String {
+        "Remove whitespace".to_string()
+    }
+}
+
+/// E204
+pub(crate) fn whitespace_after_decorator(checker: &mut Checker, decorator_list: &[Decorator]) {
+    for decorator in decorator_list {
+        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());
+
+                let start = decorator.start() + TextSize::from(1);
+                let end = start + TextSize::try_from(end).unwrap();
+                let range = TextRange::new(start, end);
+
+                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,64 @@
+---
+source: crates/ruff_linter/src/rules/pycodestyle/mod.rs
+---
+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')

scripts/check_docs_formatted.py

@@ -86,6 +86,7 @@
     "unnecessary-class-parentheses",
     "unnecessary-escaped-quote",
     "useless-semicolon",
+    "whitespace-after-decorator",
     "whitespace-after-open-bracket",
     "whitespace-before-close-bracket",
     "whitespace-before-parameters",