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,41 +255,40 @@ 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:
@@ -299,8 +296,9 @@ def bar():
 -    # 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...]]
 

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__fmtonoff_py.snap

@@ -199,7 +199,7 @@ d={'a':1,
 ```diff
 --- Black
 +++ Ruff
-@@ -1,224 +1,73 @@
+@@ -1,224 +1,72 @@
  #!/usr/bin/env python3
 -import asyncio
 -import sys
@@ -219,13 +219,11 @@ d={'a':1,
  # fmt: on
 -f"trigger 3.6 mode"
 +NOT_YET_IMPLEMENTED_ExprJoinedStr
-+
-+
  # Comment 1
  
  # Comment 2
  
--
+ 
  # fmt: off
 -def func_no_args():
 -  a; b; c
@@ -253,7 +251,12 @@ d={'a':1,
 -    offset = attr.ib(default=attr.Factory(lambda: _r.uniform(1, 2)))
 -    assert task._cancel_stack[: len(old_stack)] == old_stack
 +NOT_YET_IMPLEMENTED_StmtFunctionDef
++
++
++NOT_YET_IMPLEMENTED_StmtAsyncFunctionDef
++
  
++NOT_YET_IMPLEMENTED_StmtFunctionDef
  
 -def spaces_types(
 -    a: int = 1,
@@ -267,21 +270,21 @@ d={'a':1,
 -    i: str = r"",
 -):
 -    ...
-+NOT_YET_IMPLEMENTED_StmtAsyncFunctionDef
  
++# fmt: on
++NOT_YET_IMPLEMENTED_StmtFunctionDef
  
 -def spaces2(result=_core.Value(None)):
 -    ...
-+NOT_YET_IMPLEMENTED_StmtFunctionDef
  
++NOT_YET_IMPLEMENTED_StmtFunctionDef
  
 -something = {
 -    # fmt: off
 -    key: 'value',
 -}
-+# fmt: on
-+NOT_YET_IMPLEMENTED_StmtFunctionDef
  
++NOT_YET_IMPLEMENTED_StmtFunctionDef
  
 -def subscriptlist():
 -    atom[
@@ -292,24 +295,24 @@ d={'a':1,
 -        goes + here,
 -        andhere,
 -    ]
-+NOT_YET_IMPLEMENTED_StmtFunctionDef
  
++something = {NOT_IMPLEMENTED_dict_key: NOT_IMPLEMENTED_dict_value}
  
 -def import_as_names():
 -    # fmt: off
 -    from hello import a,        b
 -    'unformatted'
 -    # fmt: on
-+NOT_YET_IMPLEMENTED_StmtFunctionDef
  
++NOT_YET_IMPLEMENTED_StmtFunctionDef
  
 -def testlist_star_expr():
 -    # fmt: off
 -    a , b = *hello
 -    'unformatted'
 -    # fmt: on
-+something = {NOT_IMPLEMENTED_dict_key: NOT_IMPLEMENTED_dict_value}
  
++NOT_YET_IMPLEMENTED_StmtFunctionDef
  
 -def yield_expr():
 -    # fmt: off
@@ -321,8 +324,8 @@ d={'a':1,
 -    ( yield hello )
 -    'unformatted'
 -    # fmt: on
-+NOT_YET_IMPLEMENTED_StmtFunctionDef
  
++NOT_YET_IMPLEMENTED_StmtFunctionDef
  
 -def example(session):
 -    # fmt: off
@@ -333,25 +336,24 @@ d={'a':1,
 -        .order_by(models.Customer.id.asc())\
 -        .all()
 -    # fmt: on
-+NOT_YET_IMPLEMENTED_StmtFunctionDef
  
++NOT_YET_IMPLEMENTED_StmtFunctionDef
  
 -def off_and_on_without_data():
 -    """All comments here are technically on the same prefix.
-+NOT_YET_IMPLEMENTED_StmtFunctionDef
  
 -    The comments between will be formatted. This is a known limitation.
 -    """
 -    # fmt: off
++NOT_YET_IMPLEMENTED_StmtFunctionDef
  
 -    # hey, that won't work
-+NOT_YET_IMPLEMENTED_StmtFunctionDef
  
 -    # fmt: on
 -    pass
- 
 +NOT_YET_IMPLEMENTED_StmtFunctionDef
  
+ 
 -def on_and_off_broken():
 -    """Another known limitation."""
 -    # fmt: on
@@ -364,9 +366,9 @@ d={'a':1,
 -    # fmt: off
 -    # ...but comments still get reformatted even though they should not be
 -    # fmt: on
- 
 +NOT_YET_IMPLEMENTED_StmtFunctionDef
  
+ 
 -def long_lines():
 -    if True:
 -        typedargslist.extend(
@@ -406,18 +408,14 @@ d={'a':1,
 -        re.MULTILINE|re.VERBOSE
 -        # fmt: on
 -    )
- 
 +NOT_YET_IMPLEMENTED_StmtFunctionDef
  
+ 
 -def single_literal_yapf_disable():
 -    """Black does not support this."""
 -    BAZ = {(1, 2, 3, 4), (5, 6, 7, 8), (9, 10, 11, 12)}  # yapf: disable
++NOT_YET_IMPLEMENTED_StmtFunctionDef
  
-+NOT_YET_IMPLEMENTED_StmtFunctionDef
-+
-+
-+NOT_YET_IMPLEMENTED_StmtFunctionDef
-+
  
 -cfg.rule(
 -    "Default",
@@ -472,12 +470,11 @@ NOT_YET_IMPLEMENTED_StmtImportFrom
 NOT_YET_IMPLEMENTED_StmtImportFrom
 # fmt: on
 NOT_YET_IMPLEMENTED_ExprJoinedStr
-
-
 # Comment 1
 
 # Comment 2
 
+
 # fmt: off
 NOT_YET_IMPLEMENTED_StmtFunctionDef