Avoid flagging missing whitespace for decorators (#4454)

astral-sh · May 16, 2023 · d9c3f8e · d9c3f8e
1 parent 7e0d018
commit d9c3f8e
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/crates/ruff/resources/test/fixtures/pycodestyle/E22.py b/crates/ruff/resources/test/fixtures/pycodestyle/E22.py
@@ -160,6 +160,7 @@ def halves(n):
     *a, b = (1, 2, 3)
 
 
+@decorator
 def squares(n):
     return (i**2 for i in range(n))
 

diff --git a/crates/ruff/src/rules/pycodestyle/rules/logical_lines/missing_whitespace_around_operator.rs b/crates/ruff/src/rules/pycodestyle/rules/logical_lines/missing_whitespace_around_operator.rs
@@ -86,13 +86,19 @@ pub(crate) fn missing_whitespace_around_operator(
             );
 
             NeedsSpace::from(!slash_in_func)
-        } else if kind.is_unary() || kind == TokenKind::DoubleStar {
+        } else if kind.is_unary()
+            || matches!(
+                kind,
+                TokenKind::Star | TokenKind::DoubleStar | TokenKind::At
+            )
+        {
             let is_binary = prev_token.map_or(false, |prev_token| {
                 let prev_kind = prev_token.kind();
 
                 // Check if the operator is used as a binary operator.
                 // Allow unary operators: -123, -x, +1.
                 // Allow argument unpacking: foo(*args, **kwargs)
+                // Allow decorators: @foo, @foo(1)
                 matches!(
                     prev_kind,
                     TokenKind::Rpar | TokenKind::Rsqb | TokenKind::Rbrace

diff --git a/crates/ruff_python_ast/src/token_kind.rs b/crates/ruff_python_ast/src/token_kind.rs
@@ -174,7 +174,7 @@ impl TokenKind {
 
     #[inline]
     pub const fn is_unary(&self) -> bool {
-        matches!(self, TokenKind::Plus | TokenKind::Minus | TokenKind::Star)
+        matches!(self, TokenKind::Plus | TokenKind::Minus)
     }
 
     #[inline]