Format BoolOp (#4986)

crates/ruff_python_formatter/src/expression/binary_like.rs

@@ -0,0 +1,215 @@
+//! This module provides helper utilities to format an expression that has a left side, an operator,
+//! and a right side (binary like).
+
+use crate::expression::parentheses::Parentheses;
+use crate::prelude::*;
+use ruff_formatter::{format_args, write};
+use rustpython_parser::ast::Expr;
+
+/// Trait to implement a binary like syntax that has a left operand, an operator, and a right operand.
+pub(super) trait FormatBinaryLike<'ast> {
+    /// The type implementing the formatting of the operator.
+    type FormatOperator: Format<PyFormatContext<'ast>>;
+
+    /// Formats the binary like expression to `f`.
+    fn fmt_binary(
+        &self,
+        parentheses: Option<Parentheses>,
+        f: &mut PyFormatter<'ast, '_>,
+    ) -> FormatResult<()> {
+        let left = self.left()?;
+        let operator = self.operator();
+        let right = self.right()?;
+
+        let layout = if parentheses == Some(Parentheses::Custom) {
+            self.binary_layout()
+        } else {
+            BinaryLayout::Default
+        };
+
+        match layout {
+            BinaryLayout::Default => self.fmt_default(f),
+            BinaryLayout::ExpandLeft => {
+                let left = left.format().memoized();
+                let right = right.format().memoized();
+                write!(
+                    f,
+                    [best_fitting![
+                        // Everything on a single line
+                        format_args![group(&left), space(), operator, space(), right],
+                        // Break the left over multiple lines, keep the right flat
+                        format_args![
+                            group(&left).should_expand(true),
+                            space(),
+                            operator,
+                            space(),
+                            right
+                        ],
+                        // The content doesn't fit, indent the content and break before the operator.
+                        format_args![
+                            text("("),
+                            block_indent(&format_args![
+                                left,
+                                hard_line_break(),
+                                operator,
+                                space(),
+                                right
+                            ]),
+                            text(")")
+                        ]
+                    ]
+                    .with_mode(BestFittingMode::AllLines)]
+                )
+            }
+            BinaryLayout::ExpandRight => {
+                let left_group = f.group_id("BinaryLeft");
+
+                write!(
+                    f,
+                    [
+                        // Wrap the left in a group and gives it an id. The printer first breaks the
+                        // right side if `right` contains any line break because the printer breaks
+                        // sequences of groups from right to left.
+                        // Indents the left side if the group breaks.
+                        group(&format_args![
+                            if_group_breaks(&text("(")),
+                            indent_if_group_breaks(
+                                &format_args![
+                                    soft_line_break(),
+                                    left.format(),
+                                    soft_line_break_or_space(),
+                                    operator,
+                                    space()
+                                ],
+                                left_group
+                            )
+                        ])
+                        .with_group_id(Some(left_group)),
+                        // Wrap the right in a group and indents its content but only if the left side breaks
+                        group(&indent_if_group_breaks(&right.format(), left_group)),
+                        // If the left side breaks, insert a hard line break to finish the indent and close the open paren.
+                        if_group_breaks(&format_args![hard_line_break(), text(")")])
+                            .with_group_id(Some(left_group))
+                    ]
+                )
+            }
+            BinaryLayout::ExpandRightThenLeft => {
+                // The formatter expands group-sequences from right to left, and expands both if
+                // there isn't enough space when expanding only one of them.
+                write!(
+                    f,
+                    [
+                        group(&left.format()),
+                        space(),
+                        operator,
+                        space(),
+                        group(&right.format())
+                    ]
+                )
+            }
+        }
+    }
+
+    /// Determines which binary layout to use.
+    fn binary_layout(&self) -> BinaryLayout {
+        if let (Ok(left), Ok(right)) = (self.left(), self.right()) {
+            BinaryLayout::from_left_right(left, right)
+        } else {
+            BinaryLayout::Default
+        }
+    }
+
+    /// Formats the node according to the default layout.
+    fn fmt_default(&self, f: &mut PyFormatter<'ast, '_>) -> FormatResult<()>;
+
+    /// Returns the left operator
+    fn left(&self) -> FormatResult<&Expr>;
+
+    /// Returns the right operator.
+    fn right(&self) -> FormatResult<&Expr>;
+
+    /// Returns the object that formats the operator.
+    fn operator(&self) -> Self::FormatOperator;
+}
+
+fn can_break_expr(expr: &Expr) -> bool {
+    use ruff_python_ast::prelude::*;
+
+    match expr {
+        Expr::Tuple(ExprTuple {
+            elts: expressions, ..
+        })
+        | Expr::List(ExprList {
+            elts: expressions, ..
+        })
+        | Expr::Set(ExprSet {
+            elts: expressions, ..
+        })
+        | Expr::Dict(ExprDict {
+            values: expressions,
+            ..
+        }) => !expressions.is_empty(),
+        Expr::Call(ExprCall { args, keywords, .. }) => !(args.is_empty() && keywords.is_empty()),
+        Expr::ListComp(_) | Expr::SetComp(_) | Expr::DictComp(_) | Expr::GeneratorExp(_) => true,
+        Expr::UnaryOp(ExprUnaryOp { operand, .. }) => match operand.as_ref() {
+            Expr::BinOp(_) => true,
+            _ => can_break_expr(operand.as_ref()),
+        },
+        _ => false,
+    }
+}
+
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub(super) enum BinaryLayout {
+    /// Put each operand on their own line if either side expands
+    Default,
+
+    /// Try to expand the left to make it fit. Add parentheses if the left or right don't fit.
+    ///
+    ///```python
+    /// [
+    ///     a,
+    ///     b
+    /// ] & c
+    ///```
+    ExpandLeft,
+
+    /// Try to expand the right to make it fix. Add parentheses if the left or right don't fit.
+    ///
+    /// ```python
+    /// a & [
+    ///     b,
+    ///     c
+    /// ]
+    /// ```
+    ExpandRight,
+
+    /// Both the left and right side can be expanded. Try in the following order:
+    /// * expand the right side
+    /// * expand the left side
+    /// * expand both sides
+    ///
+    /// to make the expression fit
+    ///
+    /// ```python
+    /// [
+    ///     a,
+    ///     b
+    /// ] & [
+    ///     c,
+    ///     d
+    /// ]
+    /// ```
+    ExpandRightThenLeft,
+}
+
+impl BinaryLayout {
+    pub(super) fn from_left_right(left: &Expr, right: &Expr) -> BinaryLayout {
+        match (can_break_expr(left), can_break_expr(right)) {
+            (false, false) => BinaryLayout::Default,
+            (true, false) => BinaryLayout::ExpandLeft,
+            (false, true) => BinaryLayout::ExpandRight,
+            (true, true) => BinaryLayout::ExpandRightThenLeft,
+        }
+    }
+}