Format binary expressions (#4862)

* Format Binary Expressions

* Extract NeedsParentheses trait

crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/binary_expression.py

@@ -0,0 +1,43 @@
+(aaaaaaaa
+    + # trailing operator comment
+    b # trailing right comment
+)
+
+
+(aaaaaaaa # trailing left comment
+    +  # trailing operator comment
+    # leading right comment
+    b
+)
+
+
+# Black breaks the right side first for the following expressions:
+aaaaaaaaaaaaaa + caaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaal(argument1, argument2, argument3)
+aaaaaaaaaaaaaa + [bbbbbbbbbbbbbbbbbbbbbb, ccccccccccccccccccccc, dddddddddddddddd, eeeeeee]
+aaaaaaaaaaaaaa + (bbbbbbbbbbbbbbbbbbbbbb, ccccccccccccccccccccc, dddddddddddddddd, eeeeeee)
+aaaaaaaaaaaaaa + { key1:bbbbbbbbbbbbbbbbbbbbbb, key2: ccccccccccccccccccccc, key3: dddddddddddddddd, key4: eeeeeee }
+aaaaaaaaaaaaaa + { bbbbbbbbbbbbbbbbbbbbbb, ccccccccccccccccccccc, dddddddddddddddd, eeeeeee }
+aaaaaaaaaaaaaa + [a for x in bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb ]
+aaaaaaaaaaaaaa + (a for x in bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb )
+aaaaaaaaaaaaaa + {a for x in bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb}
+
+# Wraps it in parentheses if it needs to break both left and right
+aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + [
+    bbbbbbbbbbbbbbbbbbbbbb,
+    ccccccccccccccccccccc,
+    dddddddddddddddd,
+    eee
+] # comment
+
+
+
+# But only for expressions that have a statement parent.
+not (aaaaaaaaaaaaaa + {a for x in bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb})
+[a + [bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb] in c ]
+
+
+# leading comment
+(
+    # comment
+    content + b
+)

crates/ruff_python_formatter/src/builders.rs

@@ -12,7 +12,7 @@ pub(crate) trait PyFormatterExtensions<'ast, 'buf> {
     ///
     /// * [`NodeLevel::Module`]: Up to two empty lines
     /// * [`NodeLevel::CompoundStatement`]: Up to one empty line
-    /// * [`NodeLevel::Parenthesized`]: No empty lines
+    /// * [`NodeLevel::Expression`]: No empty lines
     fn join_nodes<'fmt>(&'fmt mut self, level: NodeLevel) -> JoinNodesBuilder<'fmt, 'ast, 'buf>;
 }
 
@@ -48,18 +48,18 @@ impl<'fmt, 'ast, 'buf> JoinNodesBuilder<'fmt, 'ast, 'buf> {
     {
         let node_level = self.node_level;
         let separator = format_with(|f: &mut PyFormatter| match node_level {
-            NodeLevel::TopLevel => match lines_before(f.context().contents(), node.start()) {
+            NodeLevel::TopLevel => match lines_before(node.start(), f.context().contents()) {
                 0 | 1 => hard_line_break().fmt(f),
                 2 => empty_line().fmt(f),
                 _ => write!(f, [empty_line(), empty_line()]),
             },
             NodeLevel::CompoundStatement => {
-                match lines_before(f.context().contents(), node.start()) {
+                match lines_before(node.start(), f.context().contents()) {
                     0 | 1 => hard_line_break().fmt(f),
                     _ => empty_line().fmt(f),
                 }
             }
-            NodeLevel::Parenthesized => hard_line_break().fmt(f),
+            NodeLevel::Expression => hard_line_break().fmt(f),
         });
 
         self.entry_with_separator(&separator, content);
@@ -200,7 +200,7 @@ no_leading_newline = 30"#
     // Removes all empty lines
     #[test]
     fn ranged_builder_parenthesized_level() {
-        let printed = format_ranged(NodeLevel::Parenthesized);
+        let printed = format_ranged(NodeLevel::Expression);
 
         assert_eq!(
             &printed,

crates/ruff_python_formatter/src/comments/format.rs

@@ -39,7 +39,7 @@ impl Format<PyFormatContext<'_>> for FormatLeadingComments<'_> {
         for comment in leading_comments {
             let slice = comment.slice();
 
-            let lines_after_comment = lines_after(f.context().contents(), slice.end());
+            let lines_after_comment = lines_after(slice.end(), f.context().contents());
             write!(
                 f,
                 [format_comment(comment), empty_lines(lines_after_comment)]
@@ -80,15 +80,15 @@ impl Format<PyFormatContext<'_>> for FormatLeadingAlternateBranchComments<'_> {
         if let Some(first_leading) = self.comments.first() {
             // Leading comments only preserves the lines after the comment but not before.
             // Insert the necessary lines.
-            if lines_before(f.context().contents(), first_leading.slice().start()) > 1 {
+            if lines_before(first_leading.slice().start(), f.context().contents()) > 1 {
                 write!(f, [empty_line()])?;
             }
 
             write!(f, [leading_comments(self.comments)])?;
         } else if let Some(last_preceding) = self.last_node {
             // The leading comments formatting ensures that it preserves the right amount of lines after
             // We need to take care of this ourselves, if there's no leading `else` comment.
-            if lines_after(f.context().contents(), last_preceding.end()) > 1 {
+            if lines_after(last_preceding.end(), f.context().contents()) > 1 {
                 write!(f, [empty_line()])?;
             }
         }
@@ -132,7 +132,7 @@ impl Format<PyFormatContext<'_>> for FormatTrailingComments<'_> {
             has_trailing_own_line_comment |= trailing.position().is_own_line();
 
             if has_trailing_own_line_comment {
-                let lines_before_comment = lines_before(f.context().contents(), slice.start());
+                let lines_before_comment = lines_before(slice.start(), f.context().contents());
 
                 // A trailing comment at the end of a body or list
                 // ```python
@@ -175,20 +175,26 @@ pub(crate) fn dangling_node_comments<T>(node: &T) -> FormatDanglingComments
 where
     T: AstNode,
 {
-    FormatDanglingComments {
-        node: node.as_any_node_ref(),
-    }
+    FormatDanglingComments::Node(node.as_any_node_ref())
+}
+
+pub(crate) fn dangling_comments(comments: &[SourceComment]) -> FormatDanglingComments {
+    FormatDanglingComments::Comments(comments)
 }
 
-pub(crate) struct FormatDanglingComments<'a> {
-    node: AnyNodeRef<'a>,
+pub(crate) enum FormatDanglingComments<'a> {
+    Node(AnyNodeRef<'a>),
+    Comments(&'a [SourceComment]),
 }
 
 impl Format<PyFormatContext<'_>> for FormatDanglingComments<'_> {
     fn fmt(&self, f: &mut Formatter<PyFormatContext>) -> FormatResult<()> {
         let comments = f.context().comments().clone();
 
-        let dangling_comments = comments.dangling_comments(self.node);
+        let dangling_comments = match self {
+            Self::Comments(comments) => comments,
+            Self::Node(node) => comments.dangling_comments(*node),
+        };
 
         let mut first = true;
         for comment in dangling_comments {
@@ -200,7 +206,7 @@ impl Format<PyFormatContext<'_>> for FormatDanglingComments<'_> {
                 f,
                 [
                     format_comment(comment),
-                    empty_lines(lines_after(f.context().contents(), comment.slice().end()))
+                    empty_lines(lines_after(comment.slice().end(), f.context().contents()))
                 ]
             )?;
 
@@ -301,7 +307,7 @@ impl Format<PyFormatContext<'_>> for FormatEmptyLines {
             },
 
             // Remove all whitespace in parenthesized expressions
-            NodeLevel::Parenthesized => write!(f, [hard_line_break()]),
+            NodeLevel::Expression => write!(f, [hard_line_break()]),
         }
     }
 }

crates/ruff_python_formatter/src/comments/mod.rs

@@ -103,8 +103,8 @@ use crate::comments::map::MultiMap;
 use crate::comments::node_key::NodeRefEqualityKey;
 use crate::comments::visitor::CommentsVisitor;
 pub(crate) use format::{
-    dangling_node_comments, leading_alternate_branch_comments, leading_node_comments,
-    trailing_comments, trailing_node_comments,
+    dangling_comments, dangling_node_comments, leading_alternate_branch_comments,
+    leading_node_comments, trailing_comments, trailing_node_comments,
 };
 use ruff_formatter::{SourceCode, SourceCodeSlice};
 use ruff_python_ast::node::AnyNodeRef;

crates/ruff_python_formatter/src/comments/placement.rs

@@ -520,8 +520,8 @@ fn handle_trailing_end_of_line_condition_comment<'a>(
     if preceding.ptr_eq(last_before_colon) {
         let mut start = preceding.end();
         while let Some((offset, c)) = find_first_non_trivia_character_in_range(
-            locator.contents(),
             TextRange::new(start, following.start()),
+            locator.contents(),
         ) {
             match c {
                 ':' => {
@@ -655,7 +655,7 @@ fn handle_trailing_binary_expression_left_or_operator_comment<'a>(
     );
 
     let operator_offset = loop {
-        match find_first_non_trivia_character_in_range(locator.contents(), between_operands_range) {
+        match find_first_non_trivia_character_in_range(between_operands_range, locator.contents()) {
             // Skip over closing parens
             Some((offset, ')')) => {
                 between_operands_range =
@@ -733,17 +733,17 @@ fn find_pos_only_slash_offset(
     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(
+    find_first_non_trivia_character_in_range(between_arguments_range, locator.contents()).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(),
                 ),
+                locator.contents(),
             )
             .map(|(offset, c)| {
                 debug_assert_eq!(c, '/');

crates/ruff_python_formatter/src/context.rs

@@ -82,5 +82,5 @@ pub(crate) enum NodeLevel {
     CompoundStatement,
 
     /// Formatting nodes that are enclosed in a parenthesized expression.
-    Parenthesized,
+    Expression,
 }

crates/ruff_python_formatter/src/expression/expr_attribute.rs

@@ -1,3 +1,6 @@
+use crate::expression::parentheses::{
+    default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
+};
 use crate::{verbatim_text, FormatNodeRule, PyFormatter};
 use ruff_formatter::{write, Buffer, FormatResult};
 use rustpython_parser::ast::ExprAttribute;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprAttribute> for FormatExprAttribute {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprAttribute {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}

crates/ruff_python_formatter/src/expression/expr_await.rs

@@ -1,3 +1,6 @@
+use crate::expression::parentheses::{
+    default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
+};
 use crate::{verbatim_text, FormatNodeRule, PyFormatter};
 use ruff_formatter::{write, Buffer, FormatResult};
 use rustpython_parser::ast::ExprAwait;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprAwait> for FormatExprAwait {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprAwait {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}