Re-add for functions and classes

Author: charliermarsh

crates/ruff_python_formatter/src/comments/format.rs

@@ -475,3 +475,45 @@ fn normalize_comment<'a>(
 
     Ok(Cow::Owned(std::format!("# {}", content.trim_start())))
 }
+
+/// Format the empty lines between a node and its trailing comments.
+///
+/// For example, given:
+/// ```python
+/// def func():
+///     ...
+/// # comment
+/// ```
+///
+/// This builder will insert two empty lines before the comment.
+/// ```
+pub(crate) const fn empty_lines_before_trailing_comments(
+    comments: &[SourceComment],
+    expected: u32,
+) -> FormatEmptyLinesBeforeTrailingComments {
+    FormatEmptyLinesBeforeTrailingComments { comments, expected }
+}
+
+#[derive(Copy, Clone, Debug)]
+pub(crate) struct FormatEmptyLinesBeforeTrailingComments<'a> {
+    /// The trailing comments of the node.
+    comments: &'a [SourceComment],
+    /// The expected number of empty lines before the trailing comments.
+    expected: u32,
+}
+
+impl Format<PyFormatContext<'_>> for FormatEmptyLinesBeforeTrailingComments<'_> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext>) -> FormatResult<()> {
+        if let Some(comment) = self
+            .comments
+            .iter()
+            .find(|comment| comment.line_position().is_own_line())
+        {
+            let actual = lines_before(comment.start(), f.context().source()).saturating_sub(1);
+            for _ in actual..self.expected {
+                write!(f, [empty_line()])?;
+            }
+        }
+        Ok(())
+    }
+}

crates/ruff_python_formatter/src/statement/stmt_class_def.rs

@@ -3,7 +3,9 @@ use ruff_python_ast::{Decorator, StmtClassDef};
 use ruff_python_trivia::lines_after_ignoring_trivia;
 use ruff_text_size::Ranged;
 
+use crate::comments::format::empty_lines_before_trailing_comments;
 use crate::comments::{leading_comments, trailing_comments, SourceComment};
+use crate::context::NodeLevel;
 use crate::prelude::*;
 use crate::statement::clause::{clause_body, clause_header, ClauseHeader};
 use crate::statement::suite::SuiteKind;
@@ -108,7 +110,33 @@ impl FormatNodeRule<StmtClassDef> for FormatStmtClassDef {
                 ),
                 clause_body(body, trailing_definition_comments).with_kind(SuiteKind::Class),
             ]
+        )?;
+
+        // If the class contains trailing comments, insert newlines before them.
+        // For example, given:
+        // ```python
+        // class Class:
+        //     ...
+        // # comment
+        // ```
+        //
+        // At the top-level, reformat as:
+        // ```python
+        // class Class:
+        //     ...
+        //
+        //
+        // # comment
+        // ```
+        empty_lines_before_trailing_comments(
+            comments.trailing(item),
+            if f.context().node_level() == NodeLevel::TopLevel {
+                2
+            } else {
+                1
+            },
         )
+        .fmt(f)
     }
 
     fn fmt_dangling_comments(

crates/ruff_python_formatter/src/statement/stmt_function_def.rs

@@ -1,9 +1,11 @@
+use crate::comments::format::empty_lines_before_trailing_comments;
 use ruff_formatter::write;
 use ruff_python_ast::{Parameters, StmtFunctionDef};
 use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
 use ruff_text_size::Ranged;
 
 use crate::comments::SourceComment;
+use crate::context::NodeLevel;
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::{Parentheses, Parenthesize};
 use crate::prelude::*;
@@ -144,7 +146,33 @@ impl FormatNodeRule<StmtFunctionDef> for FormatStmtFunctionDef {
                 ),
                 clause_body(body, trailing_definition_comments).with_kind(SuiteKind::Function),
             ]
+        )?;
+
+        // If the function contains trailing comments, insert newlines before them.
+        // For example, given:
+        // ```python
+        // def func():
+        //     ...
+        // # comment
+        // ```
+        //
+        // At the top-level, reformat as:
+        // ```python
+        // def func():
+        //     ...
+        //
+        //
+        // # comment
+        // ```
+        empty_lines_before_trailing_comments(
+            comments.trailing(item),
+            if f.context().node_level() == NodeLevel::TopLevel {
+                2
+            } else {
+                1
+            },
         )
+        .fmt(f)
     }
 
     fn fmt_dangling_comments(

crates/ruff_python_formatter/src/statement/suite.rs

@@ -195,12 +195,13 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
                         empty_line().fmt(f)?;
                     }
                 }
-            } else if (is_import_definition(preceding)
-                && !comments.has_trailing_own_line(preceding))
+            } else if is_import_definition(preceding)
                 && (!is_import_definition(following) || comments.has_leading(following))
             {
                 // Enforce _at least_ one empty line after an import statement (but allow up to
-                // two at the top-level).
+                // two at the top-level). In this context, "after an import statement" means that
+                // that the previous node is an import, and the following node is an import _or_ has
+                // a leading comment.
                 match self.kind {
                     SuiteKind::TopLevel => {
                         match lines_after_ignoring_trivia(preceding.end(), source) {