Place comments of left and right binary expression operands (#4751)

crates/ruff_python_formatter/src/comments/mod.rs

@@ -802,4 +802,58 @@ def test(a=10,/, # trailing positional argument comment.
 
         assert_debug_snapshot!(comments.debug(test_case.source_code));
     }
+
+    #[test]
+    fn binary_expression_left_operand_comment() {
+        let source = r#"
+a = (
+    5
+    # trailing left comment
+    +
+    # leading right comment
+    3
+)
+"#;
+        let test_case = CommentsTestCase::from_code(source);
+
+        let comments = test_case.to_comments();
+
+        assert_debug_snapshot!(comments.debug(test_case.source_code));
+    }
+
+    #[test]
+    fn binary_expression_left_operand_trailing_end_of_line_comment() {
+        let source = r#"
+a = (
+    5 # trailing left comment
+    + # trailing operator comment
+    # leading right comment
+    3
+)
+"#;
+        let test_case = CommentsTestCase::from_code(source);
+
+        let comments = test_case.to_comments();
+
+        assert_debug_snapshot!(comments.debug(test_case.source_code));
+    }
+
+    #[test]
+    fn nested_binary_expression() {
+        let source = r#"
+a = (
+    (5 # trailing left comment
+        *
+        2)
+    + # trailing operator comment
+    # leading right comment
+    3
+)
+"#;
+        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

@@ -1,6 +1,7 @@
 use crate::comments::visitor::{CommentPlacement, DecoratedComment};
 
 use crate::comments::CommentTextPosition;
+use crate::trivia::find_first_non_trivia_character_in_range;
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::source_code::Locator;
 use ruff_python_ast::whitespace;
@@ -18,6 +19,9 @@ pub(super) fn place_comment<'a>(
         .or_else(|comment| handle_in_between_bodies_comment(comment, locator))
         .or_else(|comment| handle_trailing_body_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)
+        })
 }
 
 /// Handles leading comments in front of a match case or a trailing comment of the `match` statement.
@@ -465,29 +469,135 @@ fn handle_positional_only_arguments_separator_comment<'a>(
     }
 }
 
-fn find_pos_only_slash_offset(trivia_range: TextRange, locator: &Locator) -> Option<TextSize> {
-    let mut in_comment = false;
+/// Handles comments between the left side and the operator of a binary expression (trailing comments of the left),
+/// and trailing end-of-line comments that are on the same line as the operator.
+///
+///```python
+/// a = (
+///     5 # trailing left comment
+///     + # trailing operator comment
+///     # leading right comment
+///     3
+/// )
+/// ```
+fn handle_trailing_binary_expression_left_or_operator_comment<'a>(
+    comment: DecoratedComment<'a>,
+    locator: &Locator,
+) -> CommentPlacement<'a> {
+    let Some(binary_expression) = comment.enclosing_node().expr_bin_op() else {
+        return CommentPlacement::Default(comment);
+    };
 
-    for (offset, c) in locator.slice(trivia_range).char_indices() {
-        match c {
-            '\n' | '\r' => {
-                in_comment = false;
+    // Only if there's a preceding node (in which case, the preceding node is `left`).
+    if comment.preceding_node().is_none() || comment.following_node().is_none() {
+        return CommentPlacement::Default(comment);
     }
-            '/' if !in_comment => {
-                return Some(trivia_range.start() + TextSize::try_from(offset).unwrap());
+
+    let mut between_operands_range = TextRange::new(
+        binary_expression.left.end(),
+        binary_expression.right.start(),
+    );
+
+    let operator_offset = loop {
+        match find_first_non_trivia_character_in_range(locator.contents(), between_operands_range) {
+            // Skip over closing parens
+            Some((offset, ')')) => {
+                between_operands_range =
+                    TextRange::new(offset + TextSize::new(1), between_operands_range.end());
             }
-            '#' => {
-                // SAFE because we know there's only trivia. So all content is either whitespace,
-                // or comments but never strings.
-                in_comment = true;
+            Some((offset, _)) => break offset,
+            None => return CommentPlacement::Default(comment),
         }
-            _ => {}
+    };
+
+    let comment_range = comment.slice().range();
+
+    if comment_range.end() < operator_offset {
+        // ```python
+        // a = (
+        //      5
+        //      # comment
+        //      +
+        //      3
+        // )
+        // ```
+        CommentPlacement::trailing(AnyNodeRef::from(binary_expression.left.as_ref()), comment)
+    } else if comment.text_position().is_end_of_line() {
+        // Is the operator on its own line.
+        if locator.contains_line_break(TextRange::new(
+            binary_expression.left.end(),
+            operator_offset,
+        )) && locator.contains_line_break(TextRange::new(
+            operator_offset,
+            binary_expression.right.start(),
+        )) {
+            // ```python
+            // a = (
+            //      5
+            //      + # comment
+            //      3
+            // )
+            // ```
+            CommentPlacement::dangling(binary_expression.into(), comment)
+        } else {
+            // ```python
+            // a = (
+            //      5
+            //      +
+            //      3 # comment
+            // )
+            // ```
+            // OR
+            // ```python
+            // a = (
+            //      5 # comment
+            //      +
+            //      3
+            // )
+            // ```
+            CommentPlacement::Default(comment)
+        }
+    } else {
+        // ```python
+        // a = (
+        //      5
+        //      +
+        //      # comment
+        //      3
+        // )
+        // ```
+        CommentPlacement::Default(comment)
     }
 }
 
-    None
+/// 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(
+    between_arguments_range: TextRange,
+    locator: &Locator,
+) -> Option<TextSize> {
+    // First find the comma separating the two arguments
+    find_first_non_trivia_character_in_range(locator.contents(), between_arguments_range).and_then(
+        |(comma_offset, comma)| {
+            debug_assert_eq!(comma, ',');
+
+            // Then find the position of the `/` operator
+            find_first_non_trivia_character_in_range(
+                locator.contents(),
+                TextRange::new(
+                    comma_offset + TextSize::new(1),
+                    between_arguments_range.end(),
+                ),
+            )
+            .map(|(offset, c)| {
+                debug_assert_eq!(c, '/');
+                offset
+            })
+        },
+    )
 }
 
+/// Returns `true` if `right` is `Some` and `left` and `right` are referentially equal.
 fn are_same_optional<'a, T>(left: AnyNodeRef, right: Option<T>) -> bool
 where
     T: Into<AnyNodeRef<'a>>,

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

@@ -0,0 +1,36 @@
+---
+source: crates/ruff_python_formatter/src/comments/mod.rs
+expression: comments.debug(test_case.source_code)
+---
+{
+    Node {
+        kind: ExprConstant,
+        range: 11..12,
+        source: `5`,
+    }: {
+        "leading": [],
+        "dangling": [],
+        "trailing": [
+            SourceComment {
+                text: "# trailing left comment",
+                position: OwnLine,
+                formatted: false,
+            },
+        ],
+    },
+    Node {
+        kind: ExprConstant,
+        range: 79..80,
+        source: `3`,
+    }: {
+        "leading": [
+            SourceComment {
+                text: "# leading right comment",
+                position: OwnLine,
+                formatted: false,
+            },
+        ],
+        "dangling": [],
+        "trailing": [],
+    },
+}

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

@@ -0,0 +1,51 @@
+---
+source: crates/ruff_python_formatter/src/comments/mod.rs
+expression: comments.debug(test_case.source_code)
+---
+{
+    Node {
+        kind: ExprConstant,
+        range: 11..12,
+        source: `5`,
+    }: {
+        "leading": [],
+        "dangling": [],
+        "trailing": [
+            SourceComment {
+                text: "# trailing left comment",
+                position: EndOfLine,
+                formatted: false,
+            },
+        ],
+    },
+    Node {
+        kind: ExprBinOp,
+        range: 11..104,
+        source: `5 # trailing left comment⏎`,
+    }: {
+        "leading": [],
+        "dangling": [
+            SourceComment {
+                text: "# trailing operator comment",
+                position: EndOfLine,
+                formatted: false,
+            },
+        ],
+        "trailing": [],
+    },
+    Node {
+        kind: ExprConstant,
+        range: 103..104,
+        source: `3`,
+    }: {
+        "leading": [
+            SourceComment {
+                text: "# leading right comment",
+                position: OwnLine,
+                formatted: false,
+            },
+        ],
+        "dangling": [],
+        "trailing": [],
+    },
+}

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

@@ -0,0 +1,51 @@
+---
+source: crates/ruff_python_formatter/src/comments/mod.rs
+expression: comments.debug(test_case.source_code)
+---
+{
+    Node {
+        kind: ExprBinOp,
+        range: 11..126,
+        source: `(5 # trailing left comment⏎`,
+    }: {
+        "leading": [],
+        "dangling": [
+            SourceComment {
+                text: "# trailing operator comment",
+                position: EndOfLine,
+                formatted: false,
+            },
+        ],
+        "trailing": [],
+    },
+    Node {
+        kind: ExprConstant,
+        range: 12..13,
+        source: `5`,
+    }: {
+        "leading": [],
+        "dangling": [],
+        "trailing": [
+            SourceComment {
+                text: "# trailing left comment",
+                position: EndOfLine,
+                formatted: false,
+            },
+        ],
+    },
+    Node {
+        kind: ExprConstant,
+        range: 125..126,
+        source: `3`,
+    }: {
+        "leading": [
+            SourceComment {
+                text: "# leading right comment",
+                position: OwnLine,
+                formatted: false,
+            },
+        ],
+        "dangling": [],
+        "trailing": [],
+    },
+}

crates/ruff_python_formatter/src/lib.rs

@@ -26,6 +26,7 @@ pub(crate) mod other;
 pub(crate) mod pattern;
 mod prelude;
 pub(crate) mod statement;
+mod trivia;
 
 include!("../../ruff_formatter/shared_traits.rs");
 

crates/ruff_python_formatter/src/trivia.rs

@@ -0,0 +1,41 @@
+use ruff_python_ast::whitespace::is_python_whitespace;
+use ruff_text_size::{TextLen, TextRange, TextSize};
+
+/// Searches for the first non-trivia character in `range`.
+///
+/// The search skips over any whitespace and comments.
+///
+/// Returns `Some` if the range contains any non-trivia character. The first item is the absolute offset
+/// of the character, the second item the non-trivia character.
+///
+/// Returns `None` if the range is empty or only contains trivia (whitespace or comments).
+pub(crate) fn find_first_non_trivia_character_in_range(
+    code: &str,
+    range: TextRange,
+) -> Option<(TextSize, char)> {
+    let rest = &code[range];
+    let mut char_iter = rest.chars();
+
+    while let Some(c) = char_iter.next() {
+        match c {
+            '#' => {
+                // We're now inside of a comment. Skip all content until the end of the line
+                for c in char_iter.by_ref() {
+                    if matches!(c, '\n' | '\r') {
+                        break;
+                    }
+                }
+            }
+            c => {
+                if !is_python_whitespace(c) {
+                    let index = range.start() + rest.text_len()
+                        - char_iter.as_str().text_len()
+                        - c.text_len();
+                    return Some((index, c));
+                }
+            }
+        }
+    }
+
+    None
+}