Format Function definitions (#4951)

crates/ruff_python_ast/src/function.rs

@@ -112,3 +112,22 @@ impl<'a> From<&'a StmtAsyncFunctionDef> for AnyFunctionDefinition<'a> {
         Self::AsyncFunctionDefinition(value)
     }
 }
+
+impl<'a> From<AnyFunctionDefinition<'a>> for AnyNodeRef<'a> {
+    fn from(value: AnyFunctionDefinition<'a>) -> Self {
+        match value {
+            AnyFunctionDefinition::FunctionDefinition(function_def) => {
+                AnyNodeRef::StmtFunctionDef(function_def)
+            }
+            AnyFunctionDefinition::AsyncFunctionDefinition(async_def) => {
+                AnyNodeRef::StmtAsyncFunctionDef(async_def)
+            }
+        }
+    }
+}
+
+impl<'a> From<&'a AnyFunctionDefinition<'a>> for AnyNodeRef<'a> {
+    fn from(value: &'a AnyFunctionDefinition<'a>) -> Self {
+        (*value).into()
+    }
+}

crates/ruff_python_formatter/generate.py

@@ -77,6 +77,11 @@ def to_camel_case(node: str) -> str:
     # )
     # src.joinpath(groups[group]).joinpath("mod.rs").write_text(rustfmt(mod_section))
     for node in group_nodes:
+        node_path = src.joinpath(groups[group]).joinpath(f"{to_camel_case(node)}.rs")
+        # Don't override existing manual implementations
+        if node_path.exists():
+            continue
+
         code = f"""
             use crate::{{verbatim_text, FormatNodeRule, PyFormatter}};
             use ruff_formatter::{{write, Buffer, FormatResult}};
@@ -91,9 +96,8 @@ def to_camel_case(node: str) -> str:
                 }}
             }}
             """.strip()  # noqa: E501
-        src.joinpath(groups[group]).joinpath(f"{to_camel_case(node)}.rs").write_text(
-            rustfmt(code)
-        )
+
+        node_path.write_text(rustfmt(code))
 
 # %%
 # Generate `FormatRule`, `AsFormat` and `IntoFormat`

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/function.py

@@ -0,0 +1,70 @@
+# Dangling comments
+def test(
+    # comment
+
+    # another
+
+): ...
+
+
+# Argument empty line spacing
+def test(
+    # comment
+    a,
+
+    # another
+
+    b,
+): ...
+
+
+### Different function argument wrappings
+
+def single_line(aaaaaaaaaaaaaaaaaaaaaaaaaaaaa, bbbbbbbbbbbbbbb, ccccccccccccccccc):
+    pass
+
+def arguments_on_their_own_line(aaaaaaaaaaaaaaaaaaaaaaaaaaaaa, bbbbbbbbbbbbbbb, ccccccccccc, ddddddddddddd, eeeeeee):
+    pass
+
+def argument_per_line(aaaaaaaaaaaaaaaaaaaaaaaaaaaaa, bbbbbbbbbbbbbbb, ccccccccccccccccc, ddddddddddddd, eeeeeeeeeeeeeeee, ffffffffffff):
+    pass
+
+def last_pos_only_trailing_comma(a, b, /,):
+    pass
+
+def last_pos_no_trailing_comma(a, b, /):
+    pass
+
+
+def varg_with_leading_comments(
+    a, b,
+    # comment
+    *args
+): ...
+
+def kwarg_with_leading_comments(
+    a, b,
+    # comment
+    **kwargs
+): ...
+
+def argument_with_long_default(
+    a,
+    b = ccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc + [
+        dddddddddddddddddddd, eeeeeeeeeeeeeeeeeeee, ffffffffffffffffffffffff
+    ],
+    h = []
+):  ...
+
+
+def argument_with_long_type_annotation(
+    a,
+    b: xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx | yyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyy | zzzzzzzzzzzzzzzzzzz = [0, 1, 2, 3],
+    h = []
+):  ...
+
+
+def test(): ...
+
+# Comment
+def with_leading_comment(): ...

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);

crates/ruff_python_formatter/src/expression/expr_attribute.rs

@@ -1,3 +1,4 @@
+use crate::comments::Comments;
 use crate::expression::parentheses::{
     default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
 };
@@ -39,7 +40,12 @@ impl FormatNodeRule<ExprAttribute> for FormatExprAttribute {
 }
 
 impl NeedsParentheses for ExprAttribute {
-    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
-        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    fn needs_parentheses(
+        &self,
+        parenthesize: Parenthesize,
+        source: &str,
+        comments: &Comments,
+    ) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source, comments)
     }
 }

crates/ruff_python_formatter/src/expression/expr_await.rs

@@ -1,3 +1,4 @@
+use crate::comments::Comments;
 use crate::expression::parentheses::{
     default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
 };
@@ -17,7 +18,12 @@ impl FormatNodeRule<ExprAwait> for FormatExprAwait {
 }
 
 impl NeedsParentheses for ExprAwait {
-    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
-        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    fn needs_parentheses(
+        &self,
+        parenthesize: Parenthesize,
+        source: &str,
+        comments: &Comments,
+    ) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source, comments)
     }
 }

crates/ruff_python_formatter/src/expression/expr_bin_op.rs

@@ -1,4 +1,4 @@
-use crate::comments::trailing_comments;
+use crate::comments::{trailing_comments, Comments};
 use crate::expression::parentheses::{
     default_expression_needs_parentheses, NeedsParentheses, Parenthesize,
 };
@@ -171,8 +171,13 @@ impl FormatRule<Operator, PyFormatContext<'_>> for FormatOperator {
 }
 
 impl NeedsParentheses for ExprBinOp {
-    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
-        match default_expression_needs_parentheses(self.into(), parenthesize, source) {
+    fn needs_parentheses(
+        &self,
+        parenthesize: Parenthesize,
+        source: &str,
+        comments: &Comments,
+    ) -> Parentheses {
+        match default_expression_needs_parentheses(self.into(), parenthesize, source, comments) {
             Parentheses::Optional => {
                 if should_binary_break_right_side_first(self) {
                     Parentheses::Custom

crates/ruff_python_formatter/src/expression/expr_bool_op.rs

@@ -1,3 +1,4 @@
+use crate::comments::Comments;
 use crate::expression::parentheses::{
     default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
 };
@@ -20,7 +21,12 @@ impl FormatNodeRule<ExprBoolOp> for FormatExprBoolOp {
 }
 
 impl NeedsParentheses for ExprBoolOp {
-    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
-        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    fn needs_parentheses(
+        &self,
+        parenthesize: Parenthesize,
+        source: &str,
+        comments: &Comments,
+    ) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source, comments)
     }
 }

crates/ruff_python_formatter/src/expression/expr_call.rs

@@ -1,3 +1,4 @@
+use crate::comments::Comments;
 use crate::expression::parentheses::{
     default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
 };
@@ -18,8 +19,13 @@ impl FormatNodeRule<ExprCall> for FormatExprCall {
 }
 
 impl NeedsParentheses for ExprCall {
-    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
-        match default_expression_needs_parentheses(self.into(), parenthesize, source) {
+    fn needs_parentheses(
+        &self,
+        parenthesize: Parenthesize,
+        source: &str,
+        comments: &Comments,
+    ) -> Parentheses {
+        match default_expression_needs_parentheses(self.into(), parenthesize, source, comments) {
             Parentheses::Optional => Parentheses::Never,
             parentheses => parentheses,
         }