Format Binary Expressions

`

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

@@ -0,0 +1,27 @@
+(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}
+
+
+# But only for expressions that have a statement parent.
+not (aaaaaaaaaaaaaa + {a for x in bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb})
+[a + [bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb] in c ]

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
@@ -200,7 +200,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 +301,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/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_bin_op.rs

@@ -1,12 +1,119 @@
-use crate::{verbatim_text, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
-use rustpython_parser::ast::ExprBinOp;
+use crate::comments::trailing_comments;
+use crate::prelude::*;
+use crate::FormatNodeRule;
+use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule};
+use ruff_python_ast::node::{AstNode};
+use rustpython_parser::ast::{
+ Constant, Expr, ExprAttribute, ExprBinOp, ExprConstant, ExprUnaryOp, Operator, Unaryop,
+};
 
 #[derive(Default)]
 pub struct FormatExprBinOp;
 
 impl FormatNodeRule<ExprBinOp> for FormatExprBinOp {
  fn fmt_fields(&self, item: &ExprBinOp, f: &mut PyFormatter) -> FormatResult<()> {
- write!(f, [verbatim_text(item.range)])
+ let ExprBinOp {
+ left,
+ right,
+ op,
+ range: _,
+ } = item;
+
+ let needs_space = !is_simple_power_expression(item);
+
+ let before_operator_space = if needs_space {
+ soft_line_break_or_space()
+ } else {
+ soft_line_break()
+ };
+
+ let comments = f.context().comments().clone();
+ let operator_comments = comments.dangling_comments(item.as_any_node_ref());
+
+ write!(
+ f,
+ [
+ left.format(),
+ before_operator_space,
+ op.format(),
+ trailing_comments(operator_comments),
+ ]
+ )?;
+
+ // Format the operator on its own line if it has any trailing comments at the right side has leading comments.
+ if !operator_comments.is_empty() && comments.has_leading_comments(right.as_ref().into()) {
+ write!(f, [hard_line_break()])?;
+ } else if needs_space {
+ write!(f, [space()])?;
+ }
+
+ write!(f, [group(&right.format())])
+ }
+
+ fn fmt_dangling_comments(&self, _node: &ExprBinOp, _f: &mut PyFormatter) -> FormatResult<()> {
+ // Handled inside of `fmt_fields`
+ Ok(())
+ }
+}
+
+const fn is_simple_power_expression(expr: &ExprBinOp) -> bool {
+ expr.op.is_pow() && is_simple_power_operand(&expr.left) && is_simple_power_operand(&expr.right)
+}
+
+/// Return `true` if an [`Expr`] adheres to Black's definition of a non-complex
+/// expression, in the context of a power operation.
+const fn is_simple_power_operand(expr: &Expr) -> bool {
+ match expr {
+ Expr::UnaryOp(ExprUnaryOp {
+ op: Unaryop::Not, ..
+ }) => false,
+ Expr::Constant(ExprConstant {
+ value: Constant::Complex { .. } | Constant::Float(_) | Constant::Int(_),
+ ..
+ }) => true,
+ Expr::Name(_) => true,
+ Expr::UnaryOp(ExprUnaryOp { operand, .. }) => is_simple_power_operand(operand),
+ Expr::Attribute(ExprAttribute { value, .. }) => is_simple_power_operand(value),
+ _ => false,
+ }
+}
+
+#[derive(Copy, Clone)]
+pub struct FormatOperator;
+
+impl<'ast> AsFormat<PyFormatContext<'ast>> for Operator {
+ type Format<'a> = FormatRefWithRule<'a, Operator, FormatOperator, PyFormatContext<'ast>>;
+
+ fn format(&self) -> Self::Format<'_> {
+ FormatRefWithRule::new(self, FormatOperator)
+ }
+}
+
+impl<'ast> IntoFormat<PyFormatContext<'ast>> for Operator {
+ type Format = FormatOwnedWithRule<Operator, FormatOperator, PyFormatContext<'ast>>;
+ fn into_format(self) -> Self::Format {
+ FormatOwnedWithRule::new(self, FormatOperator)
+ }
+}
+
+impl FormatRule<Operator, PyFormatContext<'_>> for FormatOperator {
+ fn fmt(&self, item: &Operator, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+ let operator = match item {
+ Operator::Add => "+",
+ Operator::Sub => "-",
+ Operator::Mult => "*",
+ Operator::MatMult => "@",
+ Operator::Div => "/",
+ Operator::Mod => "%",
+ Operator::Pow => "**",
+ Operator::LShift => "<<",
+ Operator::RShift => ">>",
+ Operator::BitOr => "|",
+ Operator::BitXor => "^",
+ Operator::BitAnd => "&",
+ Operator::FloorDiv => "//",
+ };
+
+ text(operator).fmt(f)
  }
 }