Format Function definitions (#4951)

crates/ruff_python_formatter/src/statement/stmt_async_function_def.rs

@@ -1,5 +1,6 @@
-use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
+use crate::prelude::*;
+use crate::FormatNodeRule;
+use ruff_python_ast::function::AnyFunctionDefinition;
 use rustpython_parser::ast::StmtAsyncFunctionDef;
 
 #[derive(Default)]
@@ -7,6 +8,15 @@ pub struct FormatStmtAsyncFunctionDef;
 
 impl FormatNodeRule<StmtAsyncFunctionDef> for FormatStmtAsyncFunctionDef {
     fn fmt_fields(&self, item: &StmtAsyncFunctionDef, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [not_yet_implemented(item)])
+        AnyFunctionDefinition::from(item).format().fmt(f)
+    }
+
+    fn fmt_dangling_comments(
+        &self,
+        _node: &StmtAsyncFunctionDef,
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
+        // Handled by `AnyFunctionDef`
+        Ok(())
     }
 }

crates/ruff_python_formatter/src/statement/stmt_function_def.rs

@@ -1,12 +1,141 @@
-use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
-use rustpython_parser::ast::StmtFunctionDef;
+use crate::comments::{leading_comments, trailing_comments};
+use crate::context::NodeLevel;
+use crate::expression::parentheses::Parenthesize;
+use crate::prelude::*;
+use crate::trivia::{lines_after, skip_trailing_trivia};
+use crate::FormatNodeRule;
+use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule};
+use ruff_python_ast::function::AnyFunctionDefinition;
+use rustpython_parser::ast::{Ranged, StmtFunctionDef};
 
 #[derive(Default)]
 pub struct FormatStmtFunctionDef;
 
 impl FormatNodeRule<StmtFunctionDef> for FormatStmtFunctionDef {
     fn fmt_fields(&self, item: &StmtFunctionDef, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [not_yet_implemented(item)])
+        AnyFunctionDefinition::from(item).format().fmt(f)
+    }
+
+    fn fmt_dangling_comments(
+        &self,
+        _node: &StmtFunctionDef,
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
+        // Handled by `AnyFunctionDef`
+        Ok(())
+    }
+}
+
+#[derive(Default)]
+pub struct FormatAnyFunctionDef;
+
+impl FormatRule<AnyFunctionDefinition<'_>, PyFormatContext<'_>> for FormatAnyFunctionDef {
+    fn fmt(
+        &self,
+        item: &AnyFunctionDefinition<'_>,
+        f: &mut Formatter<PyFormatContext<'_>>,
+    ) -> FormatResult<()> {
+        let comments = f.context().comments().clone();
+
+        let dangling_comments = comments.dangling_comments(item.into());
+        let trailing_definition_comments_start =
+            dangling_comments.partition_point(|comment| comment.position().is_own_line());
+
+        let (leading_function_definition_comments, trailing_definition_comments) =
+            dangling_comments.split_at(trailing_definition_comments_start);
+
+        if let Some(last_decorator) = item.decorators().last() {
+            f.join_nodes(NodeLevel::CompoundStatement)
+                .nodes(item.decorators())
+                .finish()?;
+
+            if leading_function_definition_comments.is_empty() {
+                write!(f, [hard_line_break()])?;
+            } else {
+                // Write any leading function comments (between last decorator and function header)
+                // while maintaining the right amount of empty lines between the comment
+                // and the last decorator.
+                let decorator_end =
+                    skip_trailing_trivia(last_decorator.end(), f.context().contents());
+
+                let leading_line = if lines_after(decorator_end, f.context().contents()) <= 1 {
+                    hard_line_break()
+                } else {
+                    empty_line()
+                };
+
+                write!(
+                    f,
+                    [
+                        leading_line,
+                        leading_comments(leading_function_definition_comments)
+                    ]
+                )?;
+            }
+        }
+
+        if item.is_async() {
+            write!(f, [text("async"), space()])?;
+        }
+
+        let name = item.name();
+
+        write!(
+            f,
+            [
+                text("def"),
+                space(),
+                dynamic_text(name.as_str(), None),
+                item.arguments().format(),
+            ]
+        )?;
+
+        if let Some(return_annotation) = item.returns() {
+            write!(
+                f,
+                [
+                    space(),
+                    text("->"),
+                    space(),
+                    return_annotation
+                        .format()
+                        .with_options(Parenthesize::IfBreaks)
+                ]
+            )?;
+        }
+
+        write!(
+            f,
+            [
+                text(":"),
+                trailing_comments(trailing_definition_comments),
+                block_indent(&item.body().format())
+            ]
+        )
+    }
+}
+
+impl<'def, 'ast> AsFormat<PyFormatContext<'ast>> for AnyFunctionDefinition<'def> {
+    type Format<'a> = FormatRefWithRule<
+        'a,
+        AnyFunctionDefinition<'def>,
+        FormatAnyFunctionDef,
+        PyFormatContext<'ast>,
+    > where Self: 'a;
+
+    fn format(&self) -> Self::Format<'_> {
+        FormatRefWithRule::new(self, FormatAnyFunctionDef::default())
+    }
+}
+
+impl<'def, 'ast> IntoFormat<PyFormatContext<'ast>> for AnyFunctionDefinition<'def> {
+    type Format = FormatOwnedWithRule<
+        AnyFunctionDefinition<'def>,
+        FormatAnyFunctionDef,
+        PyFormatContext<'ast>,
+    >;
+
+    fn into_format(self) -> Self::Format {
+        FormatOwnedWithRule::new(self, FormatAnyFunctionDef)
     }
 }

crates/ruff_python_formatter/src/statement/suite.rs

@@ -188,7 +188,7 @@ mod tests {
     use crate::comments::Comments;
     use crate::prelude::*;
     use crate::statement::suite::SuiteLevel;
-    use ruff_formatter::{format, SimpleFormatOptions};
+    use ruff_formatter::{format, IndentStyle, SimpleFormatOptions};
     use rustpython_parser::ast::Suite;
     use rustpython_parser::Parse;
 
@@ -216,8 +216,14 @@ def trailing_func():
 
         let statements = Suite::parse(source, "test.py").unwrap();
 
-        let context =
-            PyFormatContext::new(SimpleFormatOptions::default(), source, Comments::default());
+        let context = PyFormatContext::new(
+            SimpleFormatOptions {
+                indent_style: IndentStyle::Space(4),
+                ..SimpleFormatOptions::default()
+            },
+            source,
+            Comments::default(),
+        );
 
         let test_formatter =
             format_with(|f: &mut PyFormatter| statements.format().with_options(level).fmt(f));
@@ -252,10 +258,13 @@ NOT_YET_IMPLEMENTED_StmtClassDef
 trailing_statement = 1
 
 
-NOT_YET_IMPLEMENTED_StmtFunctionDef
+def func():
+    pass
 
 
-NOT_YET_IMPLEMENTED_StmtFunctionDef"#
+def trailing_func():
+    pass
+"#
         );
     }
 
@@ -278,9 +287,12 @@ NOT_YET_IMPLEMENTED_StmtClassDef
 
 trailing_statement = 1
 
-NOT_YET_IMPLEMENTED_StmtFunctionDef
+def func():
+    pass
 
-NOT_YET_IMPLEMENTED_StmtFunctionDef"#
+def trailing_func():
+    pass
+"#
         );
     }
 }