Handle trailing end-of-line comments in-between-bodies (#4812)

<!--
Thank you for contributing to Ruff! To help us out with reviewing, please consider the following:

- Does this pull request include a summary of the change? (See below.)
- Does this pull request include a descriptive title?
- Does this pull request include references to any relevant issues?
-->

## Summary

And more custom logic around comments in bodies... uff. 

Let's say we have the following code

```python
if x == y:
    pass # trailing comment of pass
else: # trailing comment of `else`
    print("I have no comments")
```

Right now, the formatter attaches the `# trailing comment of `else` as a trailing comment of `pass` because it doesn't "see" that there's an `else` keyword in between (because the else body is just a Vec and not a node). 

This PR adds custom logic that attaches the trailing comments after the `else` as dangling comments to the `if` statement. The if statement must then split the dangling comments by `comments.text_position()`:
* All comments up to the first end-of-line comment are leading comments of the `else` keyword.
* All end-of-line comments coming after are `trailing` comments for the `else` keyword.


## Test Plan

I added new unit tests.

crates/ruff_python_formatter/src/comments/mod.rs

@@ -872,4 +872,19 @@ a = (
 
         assert_debug_snapshot!(comments.debug(test_case.source_code));
     }
+
+    #[test]
+    fn while_trailing_else_end_of_line_comment() {
+        let source = r#"while True:
+    pass
+else: # trailing comment
+    pass
+"#;
+
+        let test_case = CommentsTestCase::from_code(source);
+
+        let comments = test_case.to_comments();
+
+        assert_debug_snapshot!(comments.debug(test_case.source_code));
+    }
 }

crates/ruff_python_formatter/src/comments/placement.rs

@@ -16,7 +16,8 @@ pub(super) fn place_comment<'a>(
 ) -> CommentPlacement<'a> {
     handle_in_between_excepthandlers_or_except_handler_and_else_or_finally_comment(comment, locator)
         .or_else(|comment| handle_match_comment(comment, locator))
-        .or_else(|comment| handle_in_between_bodies_comment(comment, locator))
+        .or_else(|comment| handle_in_between_bodies_own_line_comment(comment, locator))
+        .or_else(|comment| handle_in_between_bodies_end_of_line_comment(comment, locator))
         .or_else(|comment| handle_trailing_body_comment(comment, locator))
         .or_else(handle_trailing_end_of_line_body_comment)
         .or_else(|comment| handle_positional_only_arguments_separator_comment(comment, locator))
@@ -178,7 +179,7 @@ fn handle_in_between_excepthandlers_or_except_handler_and_else_or_finally_commen
     CommentPlacement::Default(comment)
 }
 
-/// Handles comments between the last statement and the first statement of two bodies.
+/// Handles own line comments between the last statement and the first statement of two bodies.
 ///
 /// ```python
 /// if x == y:
@@ -188,15 +189,11 @@ fn handle_in_between_excepthandlers_or_except_handler_and_else_or_finally_commen
 /// else:
 ///     print("I have no comments")
 /// ```
-fn handle_in_between_bodies_comment<'a>(
+fn handle_in_between_bodies_own_line_comment<'a>(
     comment: DecoratedComment<'a>,
     locator: &Locator,
 ) -> CommentPlacement<'a> {
-    use ruff_python_ast::prelude::*;
+    if !comment.text_position().is_own_line() {
-
-    // The rule only applies to own line comments. The default logic associates end of line comments
-    // correctly.
-    if comment.text_position().is_end_of_line() {
         return CommentPlacement::Default(comment);
     }
 
@@ -204,39 +201,7 @@ fn handle_in_between_bodies_comment<'a>(
     if let (Some(preceding), Some(following)) = (comment.preceding_node(), comment.following_node())
     {
         // ...and the following statement must be the first statement in an alternate body of the parent...
-        let is_following_the_first_statement_in_a_parents_alternate_body =
+        if !is_first_statement_in_enclosing_alternate_body(following, comment.enclosing_node()) {
-            match comment.enclosing_node() {
-                AnyNodeRef::StmtIf(StmtIf { orelse, .. })
-                | AnyNodeRef::StmtFor(StmtFor { orelse, .. })
-                | AnyNodeRef::StmtAsyncFor(StmtAsyncFor { orelse, .. })
-                | AnyNodeRef::StmtWhile(StmtWhile { orelse, .. }) => {
-                    are_same_optional(following, orelse.first())
-                }
-
-                AnyNodeRef::StmtTry(StmtTry {
-                    handlers,
-                    orelse,
-                    finalbody,
-                    ..
-                })
-                | AnyNodeRef::StmtTryStar(StmtTryStar {
-                    handlers,
-                    orelse,
-                    finalbody,
-                    ..
-                }) => {
-                    are_same_optional(following, handlers.first())
-                    // Comments between the handlers and the `else`, or comments between the `handlers` and the `finally`
-                    // are already handled by `handle_in_between_excepthandlers_or_except_handler_and_else_or_finally_comment`
-                    || handlers.is_empty() && are_same_optional(following, orelse.first())
-                    || (handlers.is_empty() || !orelse.is_empty())
-                        && are_same_optional(following, finalbody.first())
-                }
-
-                _ => false,
-            };
-
-        if !is_following_the_first_statement_in_a_parents_alternate_body {
             // ```python
             // if test:
             //     a
@@ -305,6 +270,75 @@ fn handle_in_between_bodies_comment<'a>(
     CommentPlacement::Default(comment)
 }
 
+/// Handles end of line comments comments between the last statement and the first statement of two bodies.
+///
+/// ```python
+/// if x == y:
+///     pass # trailing comment of pass
+/// else: # trailing comment of `else`
+///     print("I have no comments")
+/// ```
+fn handle_in_between_bodies_end_of_line_comment<'a>(
+    comment: DecoratedComment<'a>,
+    locator: &Locator,
+) -> CommentPlacement<'a> {
+    if !comment.text_position().is_end_of_line() {
+        return CommentPlacement::Default(comment);
+    }
+
+    // The comment must be between two statements...
+    if let (Some(preceding), Some(following)) = (comment.preceding_node(), comment.following_node())
+    {
+        // ...and the following statement must be the first statement in an alternate body of the parent...
+        if !is_first_statement_in_enclosing_alternate_body(following, comment.enclosing_node()) {
+            // ```python
+            // if test:
+            //     a
+            //     # comment
+            //     b
+            // ```
+            return CommentPlacement::Default(comment);
+        }
+
+        if !locator.contains_line_break(TextRange::new(preceding.end(), comment.slice().start())) {
+            // Trailing comment of the preceding statement
+            // ```python
+            // while test:
+            //     a # comment
+            // else:
+            //     b
+            // ```
+            CommentPlacement::trailing(preceding, comment)
+        } else if following.is_stmt_if() || following.is_except_handler() {
+            // The `elif` or except handlers have their own body to which we can attach the trailing comment
+            // ```python
+            // if test:
+            //     a
+            // elif c: # comment
+            //     b
+            // ```
+            CommentPlacement::trailing(following, comment)
+        } else {
+            // There are no bodies for the "else" branch and other bodies that are represented as a `Vec<Stmt>`.
+            // This means, there's no good place to attach the comments to.
+            // Make this a dangling comments and manually format the comment in
+            // in the enclosing node's formatting logic. For `try`, it's the formatters responsibility
+            // to correctly identify the comments for the `finally` and `orelse` block by looking
+            // at the comment's range.
+            //
+            // ```python
+            // while x == y:
+            //     pass
+            // else: # trailing
+            //     print("nooop")
+            // ```
+            CommentPlacement::dangling(comment.enclosing_node(), comment)
+        }
+    } else {
+        CommentPlacement::Default(comment)
+    }
+}
+
 /// Handles trailing comments at the end of a body block (or any other block that is indented).
 /// ```python
 /// def test():
@@ -703,3 +737,42 @@ fn last_child_in_body(node: AnyNodeRef) -> Option<AnyNodeRef> {
 
     body.last().map(AnyNodeRef::from)
 }
+
+/// Returns `true` if `following` is the first statement in an alternate `body` (e.g. the else of an if statement) of the `enclosing` node.
+fn is_first_statement_in_enclosing_alternate_body(
+    following: AnyNodeRef,
+    enclosing: AnyNodeRef,
+) -> bool {
+    use ruff_python_ast::prelude::*;
+
+    match enclosing {
+        AnyNodeRef::StmtIf(StmtIf { orelse, .. })
+        | AnyNodeRef::StmtFor(StmtFor { orelse, .. })
+        | AnyNodeRef::StmtAsyncFor(StmtAsyncFor { orelse, .. })
+        | AnyNodeRef::StmtWhile(StmtWhile { orelse, .. }) => {
+            are_same_optional(following, orelse.first())
+        }
+
+        AnyNodeRef::StmtTry(StmtTry {
+            handlers,
+            orelse,
+            finalbody,
+            ..
+        })
+        | AnyNodeRef::StmtTryStar(StmtTryStar {
+            handlers,
+            orelse,
+            finalbody,
+            ..
+        }) => {
+            are_same_optional(following, handlers.first())
+                // Comments between the handlers and the `else`, or comments between the `handlers` and the `finally`
+                // are already handled by `handle_in_between_excepthandlers_or_except_handler_and_else_or_finally_comment`
+                || handlers.is_empty() && are_same_optional(following, orelse.first())
+                || (handlers.is_empty() || !orelse.is_empty())
+                && are_same_optional(following, finalbody.first())
+        }
+
+        _ => false,
+    }
+}

crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__while_trailing_else_end_of_line_comment.snap

@@ -0,0 +1,21 @@
+---
+source: crates/ruff_python_formatter/src/comments/mod.rs
+expression: comments.debug(test_case.source_code)
+---
+{
+    Node {
+        kind: StmtWhile,
+        range: 0..54,
+        source: `while True:⏎`,
+    }: {
+        "leading": [],
+        "dangling": [
+            SourceComment {
+                text: "# trailing comment",
+                position: EndOfLine,
+                formatted: false,
+            },
+        ],
+        "trailing": [],
+    },
+}