Fix incorect placement of trailing stub function comments (#11632)

MichaReiser · web-flow · commit 9b6d2ce1f2aa · 2024-05-31T12:06:17.000Z
diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/stub_functions_trailing_comments.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/stub_functions_trailing_comments.py
@@ -0,0 +1,28 @@
+# Regression tests for https://github.com/astral-sh/ruff/issues/11569
+
+
+# comment 1
+def foo(self) -> None: ...
+def bar(self) -> None: ...
+# comment 2
+
+# comment 3
+def baz(self) -> None:
+    return None
+# comment 4
+
+
+def foo(self) -> None: ...
+# comment 5
+
+def baz(self) -> None:
+    return None
+
+
+def foo(self) -> None:
+    ... # comment 5
+def baz(self) -> None:
+    return None
+
+def foo(self) -> None: ...
+# comment 5
diff --git a/crates/ruff_python_formatter/src/comments/format.rs b/crates/ruff_python_formatter/src/comments/format.rs
@@ -164,7 +164,7 @@ impl Format<PyFormatContext<'_>> for FormatTrailingComments<'_> {
                         line_suffix(
                             &format_args![
                                 empty_lines(lines_before_comment),
-                                format_comment(trailing)
+                                format_comment(trailing),
                             ],
                             // Reserving width isn't necessary because we don't split
                             // comments and the empty lines expand any enclosing group.
@@ -535,31 +535,21 @@ fn strip_comment_prefix(comment_text: &str) -> FormatResult<&str> {
 /// ```
 ///
 /// This builder will insert a single empty line before the comment.
-pub(crate) fn empty_lines_before_trailing_comments<'a>(
-    f: &PyFormatter,
-    comments: &'a [SourceComment],
+pub(crate) fn empty_lines_before_trailing_comments(
+    comments: &[SourceComment],
     node_kind: NodeKind,
-) -> FormatEmptyLinesBeforeTrailingComments<'a> {
-    // Black has different rules for stub vs. non-stub and top level vs. indented
-    let empty_lines = match (f.options().source_type(), f.context().node_level()) {
-        (PySourceType::Stub, NodeLevel::TopLevel(_)) => 1,
-        (PySourceType::Stub, _) => u32::from(node_kind == NodeKind::StmtClassDef),
-        (_, NodeLevel::TopLevel(_)) => 2,
-        (_, _) => 1,
-    };
-
+) -> FormatEmptyLinesBeforeTrailingComments {
     FormatEmptyLinesBeforeTrailingComments {
         comments,
-        empty_lines,
+        node_kind,
     }
 }
 
 #[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.
-    empty_lines: u32,
+    node_kind: NodeKind,
 }
 
 impl Format<PyFormatContext<'_>> for FormatEmptyLinesBeforeTrailingComments<'_> {
@@ -569,9 +559,17 @@ impl Format<PyFormatContext<'_>> for FormatEmptyLinesBeforeTrailingComments<'_>
             .iter()
             .find(|comment| comment.line_position().is_own_line())
         {
+            // Black has different rules for stub vs. non-stub and top level vs. indented
+            let empty_lines = match (f.options().source_type(), f.context().node_level()) {
+                (PySourceType::Stub, NodeLevel::TopLevel(_)) => 1,
+                (PySourceType::Stub, _) => u32::from(self.node_kind == NodeKind::StmtClassDef),
+                (_, NodeLevel::TopLevel(_)) => 2,
+                (_, _) => 1,
+            };
+
             let actual = lines_before(comment.start(), f.context().source()).saturating_sub(1);
-            for _ in actual..self.empty_lines {
-                write!(f, [empty_line()])?;
+            for _ in actual..empty_lines {
+                empty_line().fmt(f)?;
             }
         }
         Ok(())
@@ -590,30 +588,16 @@ impl Format<PyFormatContext<'_>> for FormatEmptyLinesBeforeTrailingComments<'_>
 ///
 /// While `leading_comments` will preserve the existing empty line, this builder will insert an
 /// additional empty line before the comment.
-pub(crate) fn empty_lines_after_leading_comments<'a>(
-    f: &PyFormatter,
-    comments: &'a [SourceComment],
-) -> FormatEmptyLinesAfterLeadingComments<'a> {
-    // Black has different rules for stub vs. non-stub and top level vs. indented
-    let empty_lines = match (f.options().source_type(), f.context().node_level()) {
-        (PySourceType::Stub, NodeLevel::TopLevel(_)) => 1,
-        (PySourceType::Stub, _) => 0,
-        (_, NodeLevel::TopLevel(_)) => 2,
-        (_, _) => 1,
-    };
-
-    FormatEmptyLinesAfterLeadingComments {
-        comments,
-        empty_lines,
-    }
+pub(crate) fn empty_lines_after_leading_comments(
+    comments: &[SourceComment],
+) -> FormatEmptyLinesAfterLeadingComments {
+    FormatEmptyLinesAfterLeadingComments { comments }
 }
 
 #[derive(Copy, Clone, Debug)]
 pub(crate) struct FormatEmptyLinesAfterLeadingComments<'a> {
     /// The leading comments of the node.
     comments: &'a [SourceComment],
-    /// The expected number of empty lines after the leading comments.
-    empty_lines: u32,
 }
 
 impl Format<PyFormatContext<'_>> for FormatEmptyLinesAfterLeadingComments<'_> {
@@ -624,6 +608,14 @@ impl Format<PyFormatContext<'_>> for FormatEmptyLinesAfterLeadingComments<'_> {
             .rev()
             .find(|comment| comment.line_position().is_own_line())
         {
+            // Black has different rules for stub vs. non-stub and top level vs. indented
+            let empty_lines = match (f.options().source_type(), f.context().node_level()) {
+                (PySourceType::Stub, NodeLevel::TopLevel(_)) => 1,
+                (PySourceType::Stub, _) => 0,
+                (_, NodeLevel::TopLevel(_)) => 2,
+                (_, _) => 1,
+            };
+
             let actual = lines_after(comment.end(), f.context().source()).saturating_sub(1);
             // If there are no empty lines, keep the comment tight to the node.
             if actual == 0 {
@@ -632,12 +624,12 @@ impl Format<PyFormatContext<'_>> for FormatEmptyLinesAfterLeadingComments<'_> {
 
             // If there are more than enough empty lines already, `leading_comments` will
             // trim them as necessary.
-            if actual >= self.empty_lines {
+            if actual >= empty_lines {
                 return Ok(());
             }
 
-            for _ in actual..self.empty_lines {
-                write!(f, [empty_line()])?;
+            for _ in actual..empty_lines {
+                empty_line().fmt(f)?;
             }
         }
         Ok(())
diff --git a/crates/ruff_python_formatter/src/statement/stmt_class_def.rs b/crates/ruff_python_formatter/src/statement/stmt_class_def.rs
@@ -55,7 +55,7 @@ impl FormatNodeRule<StmtClassDef> for FormatStmtClassDef {
         // newline between the comment and the node, but we _require_ two newlines. If there are
         // _no_ newlines between the comment and the node, we don't insert _any_ newlines; if there
         // are more than two, then `leading_comments` will preserve the correct number of newlines.
-        empty_lines_after_leading_comments(f, comments.leading(item)).fmt(f)?;
+        empty_lines_after_leading_comments(comments.leading(item)).fmt(f)?;
 
         write!(
             f,
@@ -152,7 +152,7 @@ impl FormatNodeRule<StmtClassDef> for FormatStmtClassDef {
         //
         // # comment
         // ```
-        empty_lines_before_trailing_comments(f, comments.trailing(item), NodeKind::StmtClassDef)
+        empty_lines_before_trailing_comments(comments.trailing(item), NodeKind::StmtClassDef)
             .fmt(f)?;
 
         Ok(())
diff --git a/crates/ruff_python_formatter/src/statement/stmt_function_def.rs b/crates/ruff_python_formatter/src/statement/stmt_function_def.rs
@@ -52,7 +52,7 @@ impl FormatNodeRule<StmtFunctionDef> for FormatStmtFunctionDef {
         // newline between the comment and the node, but we _require_ two newlines. If there are
         // _no_ newlines between the comment and the node, we don't insert _any_ newlines; if there
         // are more than two, then `leading_comments` will preserve the correct number of newlines.
-        empty_lines_after_leading_comments(f, comments.leading(item)).fmt(f)?;
+        empty_lines_after_leading_comments(comments.leading(item)).fmt(f)?;
 
         write!(
             f,
@@ -86,7 +86,7 @@ impl FormatNodeRule<StmtFunctionDef> for FormatStmtFunctionDef {
         //
         // # comment
         // ```
-        empty_lines_before_trailing_comments(f, comments.trailing(item), NodeKind::StmtFunctionDef)
+        empty_lines_before_trailing_comments(comments.trailing(item), NodeKind::StmtFunctionDef)
             .fmt(f)
     }
 }
diff --git a/crates/ruff_python_formatter/src/statement/suite.rs b/crates/ruff_python_formatter/src/statement/suite.rs
@@ -240,7 +240,8 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
                                     preceding_stub.end(),
                                     f.context().source(),
                                 ) < 2
-                            });
+                            })
+                        && !preceding_comments.has_trailing_own_line();
 
                     if !is_preceding_stub_function_without_empty_line {
                         match self.kind {
diff --git a/crates/ruff_python_formatter/tests/snapshots/format@statement__stub_functions_trailing_comments.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@statement__stub_functions_trailing_comments.py.snap
@@ -0,0 +1,77 @@
+---
+source: crates/ruff_python_formatter/tests/fixtures.rs
+input_file: crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/stub_functions_trailing_comments.py
+---
+## Input
+```python
+# Regression tests for https://github.com/astral-sh/ruff/issues/11569
+
+
+# comment 1
+def foo(self) -> None: ...
+def bar(self) -> None: ...
+# comment 2
+
+# comment 3
+def baz(self) -> None:
+    return None
+# comment 4
+
+
+def foo(self) -> None: ...
+# comment 5
+
+def baz(self) -> None:
+    return None
+
+
+def foo(self) -> None:
+    ... # comment 5
+def baz(self) -> None:
+    return None
+
+def foo(self) -> None: ...
+# comment 5
+```
+
+## Output
+```python
+# Regression tests for https://github.com/astral-sh/ruff/issues/11569
+
+
+# comment 1
+def foo(self) -> None: ...
+def bar(self) -> None: ...
+
+
+# comment 2
+
+
+# comment 3
+def baz(self) -> None:
+    return None
+
+
+# comment 4
+
+
+def foo(self) -> None: ...
+
+
+# comment 5
+
+
+def baz(self) -> None:
+    return None
+
+
+def foo(self) -> None: ...  # comment 5
+def baz(self) -> None:
+    return None
+
+
+def foo(self) -> None: ...
+
+
+# comment 5
+```
