Trailing own line comments before func or class (#4921)

crates/ruff_python_formatter/src/comments/placement.rs

@@ -1,7 +1,7 @@
 use crate::comments::visitor::{CommentPlacement, DecoratedComment};
-
 use crate::comments::CommentTextPosition;
 use crate::trivia::find_first_non_trivia_character_in_range;
+use ruff_newlines::StrExt;
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::source_code::Locator;
 use ruff_python_ast::whitespace;
@@ -21,6 +21,9 @@ pub(super) fn place_comment<'a>(
         .or_else(|comment| handle_trailing_body_comment(comment, locator))
         .or_else(handle_trailing_end_of_line_body_comment)
         .or_else(|comment| handle_trailing_end_of_line_condition_comment(comment, locator))
+        .or_else(|comment| {
+            handle_module_level_own_line_comment_before_class_or_function_comment(comment, locator)
+        })
         .or_else(|comment| handle_positional_only_arguments_separator_comment(comment, locator))
         .or_else(|comment| {
             handle_trailing_binary_expression_left_or_operator_comment(comment, locator)
@@ -726,6 +729,76 @@ fn handle_trailing_binary_expression_left_or_operator_comment<'a>(
     }
 }
 
+/// Handles own line comments on the module level before a class or function statement.
+/// A comment only becomes the leading comment of a class or function if it isn't separated by an empty
+/// line from the class. Comments that are separated by at least one empty line from the header of the
+/// class are considered trailing comments of the previous statement.
+///
+/// This handling is necessary because Ruff inserts two empty lines before each class or function.
+/// Let's take this example:
+///
+/// ```python
+/// some = statement
+/// # This should be stick to the statement above
+///
+///
+/// # This should be split from the above by two lines
+/// class MyClassWithComplexLeadingComments:
+///     pass
+/// ```
+///
+/// By default, the `# This should be stick to the statement above` would become a leading comment
+/// of the `class` AND the `Suite` formatting separates the comment by two empty lines from the
+/// previous statement, so that the result becomes:
+///
+/// ```python
+/// some = statement
+///
+///
+/// # This should be stick to the statement above
+///
+///
+/// # This should be split from the above by two lines
+/// class MyClassWithComplexLeadingComments:
+///     pass
+/// ```
+///
+/// Which is not what we want. The work around is to make the `# This should be stick to the statement above`
+/// a trailing comment of the previous statement.
+fn handle_module_level_own_line_comment_before_class_or_function_comment<'a>(
+    comment: DecoratedComment<'a>,
+    locator: &Locator,
+) -> CommentPlacement<'a> {
+    // Only applies for own line comments on the module level...
+    if !comment.text_position().is_own_line() || !comment.enclosing_node().is_module() {
+        return CommentPlacement::Default(comment);
+    }
+
+    // ... for comments with a preceding and following node,
+    let (Some(preceding), Some(following)) = (comment.preceding_node(), comment.following_node()) else {
+        return CommentPlacement::Default(comment)
+    };
+
+    // ... where the following is a function or class statement.
+    if !matches!(
+        following,
+        AnyNodeRef::StmtAsyncFunctionDef(_)
+            | AnyNodeRef::StmtFunctionDef(_)
+            | AnyNodeRef::StmtClassDef(_)
+    ) {
+        return CommentPlacement::Default(comment);
+    }
+
+    // Make the comment a leading comment if there's no empty line between the comment and the function / class header
+    if max_empty_lines(locator.slice(TextRange::new(comment.slice().end(), following.start()))) == 0
+    {
+        CommentPlacement::leading(following, comment)
+    } else {
+        // Otherwise attach the comment as trailing comment to the previous statement
+        CommentPlacement::trailing(preceding, comment)
+    }
+}
+
 /// Finds the offset of the `/` that separates the positional only and arguments from the other arguments.
 /// Returns `None` if the positional only separator `/` isn't present in the specified range.
 fn find_pos_only_slash_offset(
@@ -866,3 +939,70 @@ fn is_first_statement_in_enclosing_alternate_body(
         _ => false,
     }
 }
+
+/// Counts the number of newlines in `contents`.
+fn max_empty_lines(contents: &str) -> usize {
+    let mut empty_lines = 0;
+    let mut max_empty_lines = 0;
+
+    for line in contents.universal_newlines().skip(1) {
+        if line.trim().is_empty() {
+            empty_lines += 1;
+        } else {
+            max_empty_lines = max_empty_lines.max(empty_lines);
+            empty_lines = 0;
+        }
+    }
+
+    max_empty_lines
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::comments::placement::max_empty_lines;
+
+    #[test]
+    fn count_empty_lines_in_trivia() {
+        assert_eq!(max_empty_lines(""), 0);
+        assert_eq!(max_empty_lines("# trailing comment\n # other comment\n"), 0);
+        assert_eq!(
+            max_empty_lines("# trailing comment\n# own line comment\n"),
+            0
+        );
+        assert_eq!(
+            max_empty_lines("# trailing comment\n\n# own line comment\n"),
+            1
+        );
+
+        assert_eq!(
+            max_empty_lines(
+                "# trailing comment\n\n# own line comment\n\n# an other own line comment"
+            ),
+            1
+        );
+
+        assert_eq!(
+            max_empty_lines(
+                "# trailing comment\n\n# own line comment\n\n# an other own line comment\n# block"
+            ),
+            1
+        );
+
+        assert_eq!(
+            max_empty_lines(
+                "# trailing comment\n\n# own line comment\n\n\n# an other own line comment\n# block"
+            ),
+            2
+        );
+
+        assert_eq!(
+            max_empty_lines(
+                r#"# This multiline comments section
+# should be split from the statement
+# above by two lines.
+"#
+            ),
+            0
+        );
+    }
+}

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments9_py.snap

@@ -162,7 +162,7 @@ def bar():
  
  
  some = statement
-@@ -14,148 +13,78 @@
+@@ -14,24 +13,21 @@
  # This multiline comments section
  # should be split from the statement
  # above by two lines.
@@ -190,11 +190,9 @@ def bar():
  
  
  some = statement
-+
-+
- # This should be stick to the statement above
+@@ -39,123 +35,55 @@
+ 
  
--
  # This should be split from the above by two lines
 -class MyClassWithComplexLeadingComments:
 -    pass
@@ -257,50 +255,50 @@ def bar():
 -# leading 4 that already has an empty line
 -def decorated_with_split_leading_comments():
 -    pass
--
--
++NOT_YET_IMPLEMENTED_StmtFunctionDef
+ 
+ 
 -def main():
 -    if a:
 -        # Leading comment before inline function
 -        def inline():
 -            pass
-+NOT_YET_IMPLEMENTED_StmtFunctionDef
- 
+-
 -        # Another leading comment
 -        def another_inline():
 -            pass
++NOT_YET_IMPLEMENTED_StmtFunctionDef
  
 -    else:
 -        # More leading comments
 -        def inline_after_else():
 -            pass
-+NOT_YET_IMPLEMENTED_StmtFunctionDef
- 
+-
  
 -if a:
 -    # Leading comment before "top-level inline" function
 -    def top_level_quote_inline():
 -        pass
-+NOT_YET_IMPLEMENTED_StmtIf
- 
+-
 -    # Another leading comment
 -    def another_top_level_quote_inline_inline():
 -        pass
- 
+-
 -else:
 -    # More leading comments
 -    def top_level_quote_inline_after_else():
 -        pass
-+NOT_YET_IMPLEMENTED_StmtClassDef
++NOT_YET_IMPLEMENTED_StmtIf
  
  
 -class MyClass:
 -    # First method has no empty lines between bare class def.
 -    # More comments.
 -    def first_method(self):
 -        pass
--
--
++NOT_YET_IMPLEMENTED_StmtClassDef
+ 
+ 
  # Regression test for https://github.com/psf/black/issues/3454.
 -def foo():
 -    pass
@@ -367,10 +365,9 @@ NOT_YET_IMPLEMENTED_StmtClassDef
 
 
 some = statement
-
-
 # This should be stick to the statement above
 
+
 # This should be split from the above by two lines
 NOT_YET_IMPLEMENTED_StmtClassDef
 

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__fmtonoff2_py.snap

@@ -60,9 +60,9 @@ def test_calculate_fades():
  TmSt = 1
  TmEx = 2
  
-+
  # fmt: off
  
++
  # Test data:
  #   Position, Volume, State, TmSt/TmEx/None, [call, [arg1...]]
  
@@ -113,9 +113,9 @@ NOT_YET_IMPLEMENTED_StmtImport
 TmSt = 1
 TmEx = 2
 
-
 # fmt: off
 
+
 # Test data:
 #   Position, Volume, State, TmSt/TmEx/None, [call, [arg1...]]