Format Function definitions (#4951)

crates/ruff_python_formatter/src/comments/format.rs

@@ -1,7 +1,7 @@
 use crate::comments::SourceComment;
 use crate::context::NodeLevel;
 use crate::prelude::*;
-use crate::trivia::{lines_after, lines_before};
+use crate::trivia::{lines_after, lines_before, skip_trailing_trivia};
 use ruff_formatter::{format_args, write, FormatError, SourceCode};
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::prelude::AstNode;
@@ -86,9 +86,10 @@ impl Format<PyFormatContext<'_>> for FormatLeadingAlternateBranchComments<'_> {
 
             write!(f, [leading_comments(self.comments)])?;
         } else if let Some(last_preceding) = self.last_node {
+            let full_end = skip_trailing_trivia(last_preceding.end(), f.context().contents());
             // The leading comments formatting ensures that it preserves the right amount of lines after
             // We need to take care of this ourselves, if there's no leading `else` comment.
-            if lines_after(last_preceding.end(), f.context().contents()) > 1 {
+            if lines_after(full_end, f.context().contents()) > 1 {
                 write!(f, [empty_line()])?;
             }
         }

crates/ruff_python_formatter/src/comments/mod.rs

@@ -103,7 +103,7 @@ use crate::comments::map::MultiMap;
 use crate::comments::node_key::NodeRefEqualityKey;
 use crate::comments::visitor::CommentsVisitor;
 pub(crate) use format::{
-    dangling_comments, dangling_node_comments, leading_alternate_branch_comments,
+    dangling_comments, dangling_node_comments, leading_alternate_branch_comments, leading_comments,
     leading_node_comments, trailing_comments, trailing_node_comments,
 };
 use ruff_formatter::{SourceCode, SourceCodeSlice};
@@ -295,6 +295,14 @@ impl<'a> Comments<'a> {
         !self.trailing_comments(node).is_empty()
     }
 
+    /// Returns `true` if the given `node` has any [trailing own line comments](self#trailing-comments).
+    #[inline]
+    pub(crate) fn has_trailing_own_line_comments(&self, node: AnyNodeRef) -> bool {
+        self.trailing_comments(node)
+            .iter()
+            .any(|comment| comment.position().is_own_line())
+    }
+
     /// Returns an iterator over the [leading](self#leading-comments) and [trailing comments](self#trailing-comments) of `node`.
     pub(crate) fn leading_trailing_comments(
         &self,

crates/ruff_python_formatter/src/comments/placement.rs

@@ -28,6 +28,7 @@ pub(super) fn place_comment<'a>(
         .or_else(|comment| {
             handle_trailing_binary_expression_left_or_operator_comment(comment, locator)
         })
+        .or_else(handle_leading_function_with_decorators_comment)
 }
 
 /// Handles leading comments in front of a match case or a trailing comment of the `match` statement.
@@ -103,9 +104,8 @@ fn handle_match_comment<'a>(
         }
     } else {
         // Comment after the last statement in a match case...
-        let match_stmt_indentation = whitespace::indentation(locator, match_stmt)
-            .unwrap_or_default()
-            .len();
+        let match_stmt_indentation =
+            whitespace::indentation(locator, match_stmt).map_or(usize::MAX, str::len);
 
         if comment_indentation <= match_case_indentation
             && comment_indentation > match_stmt_indentation
@@ -379,9 +379,8 @@ fn handle_trailing_body_comment<'a>(
     let mut grand_parent_body = None;
 
     loop {
-        let child_indentation = whitespace::indentation(locator, &current_child)
-            .map(str::len)
-            .unwrap_or_default();
+        let child_indentation =
+            whitespace::indentation(locator, &current_child).map_or(usize::MAX, str::len);
 
         match comment_indentation_len.cmp(&child_indentation) {
             Ordering::Less => {
@@ -511,6 +510,11 @@ fn handle_trailing_end_of_line_condition_comment<'a>(
         | AnyNodeRef::StmtAsyncWith(StmtAsyncWith { items, .. }) => {
             items.last().map(AnyNodeRef::from)
         }
+        AnyNodeRef::StmtFunctionDef(StmtFunctionDef { returns, args, .. })
+        | AnyNodeRef::StmtAsyncFunctionDef(StmtAsyncFunctionDef { returns, args, .. }) => returns
+            .as_deref()
+            .map(AnyNodeRef::from)
+            .or_else(|| Some(AnyNodeRef::from(args.as_ref()))),
         _ => None,
     };
 
@@ -820,6 +824,32 @@ fn find_pos_only_slash_offset(
     None
 }
 
+/// Handles own line comments between the last function decorator and the *header* of the function.
+/// It attaches these comments as dangling comments to the function instead of making them
+/// leading argument comments.
+///
+/// ```python
+/// @decorator
+/// # leading function comment
+/// def test():
+///      ...
+/// ```
+fn handle_leading_function_with_decorators_comment(comment: DecoratedComment) -> CommentPlacement {
+    let is_preceding_decorator = comment
+        .preceding_node()
+        .map_or(false, |node| node.is_decorator());
+
+    let is_following_arguments = comment
+        .following_node()
+        .map_or(false, |node| node.is_arguments());
+
+    if comment.text_position().is_own_line() && is_preceding_decorator && is_following_arguments {
+        CommentPlacement::dangling(comment.enclosing_node(), comment)
+    } else {
+        CommentPlacement::Default(comment)
+    }
+}
+
 /// Returns `true` if `right` is `Some` and `left` and `right` are referentially equal.
 fn are_same_optional<'a, T>(left: AnyNodeRef, right: Option<T>) -> bool
 where

crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__trailing_comment_after_single_statement_body.snap

@@ -4,18 +4,18 @@ expression: comments.debug(test_case.source_code)
 ---
 {
     Node {
-        kind: StmtPass,
-        range: 12..16,
-        source: `pass`,
+        kind: StmtExpr,
+        range: 29..42,
+        source: `print("test")`,
     }: {
-        "leading": [],
-        "dangling": [],
-        "trailing": [
+        "leading": [
             SourceComment {
                 text: "# Test",
                 position: OwnLine,
                 formatted: false,
             },
         ],
+        "dangling": [],
+        "trailing": [],
     },
 }

crates/ruff_python_formatter/src/comments/visitor.rs

@@ -165,6 +165,13 @@ impl<'ast> PreorderVisitor<'ast> for CommentsVisitor<'ast> {
         self.finish_node(expr);
     }
 
+    fn visit_decorator(&mut self, decorator: &'ast Decorator) {
+        if self.start_node(decorator).is_traverse() {
+            walk_decorator(self, decorator);
+        }
+        self.finish_node(decorator);
+    }
+
     fn visit_expr(&mut self, expr: &'ast Expr) {
         if self.start_node(expr).is_traverse() {
             walk_expr(self, expr);