Format binary expressions (#4862)

* Format Binary Expressions

* Extract NeedsParentheses trait

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)
+    }
+}

crates/ruff_python_formatter/src/expression/expr_bin_op.rs

@@ -1,12 +1,215 @@
-use crate::{verbatim_text, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
-use rustpython_parser::ast::ExprBinOp;
+use crate::comments::trailing_comments;
+use crate::expression::parentheses::{
+    default_expression_needs_parentheses, NeedsParentheses, Parenthesize,
+};
+use crate::expression::Parentheses;
+use crate::prelude::*;
+use crate::FormatNodeRule;
+use ruff_formatter::{
+    format_args, write, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions,
+};
+use ruff_python_ast::node::AstNode;
+use rustpython_parser::ast::{
+    Constant, Expr, ExprAttribute, ExprBinOp, ExprConstant, ExprUnaryOp, Operator, Unaryop,
+};
 
 #[derive(Default)]
-pub struct FormatExprBinOp;
+pub struct FormatExprBinOp {
+    parentheses: Option<Parentheses>,
+}
+
+impl FormatRuleWithOptions<ExprBinOp, PyFormatContext<'_>> for FormatExprBinOp {
+    type Options = Option<Parentheses>;
+
+    fn with_options(mut self, options: Self::Options) -> Self {
+        self.parentheses = options;
+        self
+    }
+}
 
 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 should_break_right = self.parentheses == Some(Parentheses::Custom);
+
+        if should_break_right {
+            let left = left.format().memoized();
+            let right = right.format().memoized();
+
+            write!(
+                f,
+                [best_fitting![
+                    // The whole expression on a single line
+                    format_args![left, space(), op.format(), space(), right],
+                    // Break the right, but keep the left flat
+                    format_args![
+                        left,
+                        space(),
+                        op.format(),
+                        space(),
+                        group(&right).should_expand(true),
+                    ],
+                    // Break after the operator, try to keep the right flat, otherwise expand it
+                    format_args![
+                        text("("),
+                        block_indent(&format_args![
+                            left,
+                            hard_line_break(),
+                            op.format(),
+                            space(),
+                            group(&right),
+                        ]),
+                        text(")")
+                    ],
+                ]]
+            )
+        } else {
+            let comments = f.context().comments().clone();
+            let operator_comments = comments.dangling_comments(item.as_any_node_ref());
+            let needs_space = !is_simple_power_expression(item);
+
+            let before_operator_space = if needs_space {
+                soft_line_break_or_space()
+            } else {
+                soft_line_break()
+            };
+
+            write!(
+                f,
+                [
+                    left.format(),
+                    before_operator_space,
+                    op.format(),
+                    trailing_comments(operator_comments),
+                ]
+            )?;
+
+            // Format the operator on its own line if the operator has trailing comments and 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](https://black.readthedocs.io/en/stable/the_black_code_style/current_style.html#line-breaks-binary-operators)
+/// 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)
+    }
+}
+
+impl NeedsParentheses for ExprBinOp {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        match default_expression_needs_parentheses(self.into(), parenthesize, source) {
+            Parentheses::Optional => {
+                if should_binary_break_right_side_first(self) {
+                    Parentheses::Custom
+                } else {
+                    Parentheses::Optional
+                }
+            }
+            parentheses => parentheses,
+        }
+    }
+}
+
+pub(super) fn should_binary_break_right_side_first(expr: &ExprBinOp) -> bool {
+    use ruff_python_ast::prelude::*;
+
+    if expr.left.is_bin_op_expr() {
+        false
+    } else {
+        match expr.right.as_ref() {
+            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
+            }
+            _ => false,
+        }
     }
 }

crates/ruff_python_formatter/src/expression/expr_bool_op.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::ExprBoolOp;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprBoolOp> for FormatExprBoolOp {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprBoolOp {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}

crates/ruff_python_formatter/src/expression/expr_call.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::ExprCall;
@@ -10,3 +13,12 @@ impl FormatNodeRule<ExprCall> for FormatExprCall {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprCall {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        match default_expression_needs_parentheses(self.into(), parenthesize, source) {
+            Parentheses::Optional => Parentheses::Never,
+            parentheses => parentheses,
+        }
+    }
+}

crates/ruff_python_formatter/src/expression/expr_compare.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::ExprCompare;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprCompare> for FormatExprCompare {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprCompare {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}

crates/ruff_python_formatter/src/expression/expr_constant.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::ExprConstant;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprConstant> for FormatExprConstant {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprConstant {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}

crates/ruff_python_formatter/src/expression/expr_dict.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::ExprDict;
@@ -10,3 +13,12 @@ impl FormatNodeRule<ExprDict> for FormatExprDict {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprDict {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        match default_expression_needs_parentheses(self.into(), parenthesize, source) {
+            Parentheses::Optional => Parentheses::Never,
+            parentheses => parentheses,
+        }
+    }
+}

crates/ruff_python_formatter/src/expression/expr_dict_comp.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::ExprDictComp;
@@ -10,3 +13,12 @@ impl FormatNodeRule<ExprDictComp> for FormatExprDictComp {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprDictComp {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        match default_expression_needs_parentheses(self.into(), parenthesize, source) {
+            Parentheses::Optional => Parentheses::Never,
+            parentheses => parentheses,
+        }
+    }
+}

crates/ruff_python_formatter/src/expression/expr_formatted_value.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::ExprFormattedValue;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprFormattedValue> for FormatExprFormattedValue {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprFormattedValue {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}

crates/ruff_python_formatter/src/expression/expr_generator_exp.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::ExprGeneratorExp;
@@ -10,3 +13,12 @@ impl FormatNodeRule<ExprGeneratorExp> for FormatExprGeneratorExp {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprGeneratorExp {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        match default_expression_needs_parentheses(self.into(), parenthesize, source) {
+            Parentheses::Optional => Parentheses::Never,
+            parentheses => parentheses,
+        }
+    }
+}

crates/ruff_python_formatter/src/expression/expr_if_exp.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::ExprIfExp;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprIfExp> for FormatExprIfExp {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprIfExp {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}

crates/ruff_python_formatter/src/expression/expr_joined_str.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::ExprJoinedStr;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprJoinedStr> for FormatExprJoinedStr {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprJoinedStr {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}

crates/ruff_python_formatter/src/expression/expr_lambda.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::ExprLambda;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprLambda> for FormatExprLambda {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprLambda {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}

crates/ruff_python_formatter/src/expression/expr_list.rs

@@ -1,5 +1,10 @@
-use crate::{verbatim_text, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
+use crate::comments::dangling_comments;
+use crate::expression::parentheses::{
+    default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
+};
+use crate::prelude::*;
+use crate::FormatNodeRule;
+use ruff_formatter::{format_args, write};
 use rustpython_parser::ast::ExprList;
 
 #[derive(Default)]
@@ -7,6 +12,55 @@ pub struct FormatExprList;
 
 impl FormatNodeRule<ExprList> for FormatExprList {
     fn fmt_fields(&self, item: &ExprList, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [verbatim_text(item.range)])
+        let ExprList {
+            range: _,
+            elts,
+            ctx: _,
+        } = item;
+
+        let items = format_with(|f| {
+            let mut iter = elts.iter();
+
+            if let Some(first) = iter.next() {
+                write!(f, [first.format()])?;
+            }
+
+            for item in iter {
+                write!(f, [text(","), soft_line_break_or_space(), item.format()])?;
+            }
+
+            if !elts.is_empty() {
+                write!(f, [if_group_breaks(&text(","))])?;
+            }
+
+            Ok(())
+        });
+
+        let comments = f.context().comments().clone();
+        let dangling = comments.dangling_comments(item.into());
+
+        write!(
+            f,
+            [group(&format_args![
+                text("["),
+                dangling_comments(dangling),
+                soft_block_indent(&items),
+                text("]")
+            ])]
+        )
+    }
+
+    fn fmt_dangling_comments(&self, _node: &ExprList, _f: &mut PyFormatter) -> FormatResult<()> {
+        // Handled as part of `fmt_fields`
+        Ok(())
+    }
+}
+
+impl NeedsParentheses for ExprList {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        match default_expression_needs_parentheses(self.into(), parenthesize, source) {
+            Parentheses::Optional => Parentheses::Never,
+            parentheses => parentheses,
+        }
     }
 }

crates/ruff_python_formatter/src/expression/expr_list_comp.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::ExprListComp;
@@ -10,3 +13,12 @@ impl FormatNodeRule<ExprListComp> for FormatExprListComp {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprListComp {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        match default_expression_needs_parentheses(self.into(), parenthesize, source) {
+            Parentheses::Optional => Parentheses::Never,
+            parentheses => parentheses,
+        }
+    }
+}

crates/ruff_python_formatter/src/expression/expr_name.rs

@@ -1,3 +1,6 @@
+use crate::expression::parentheses::{
+    default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
+};
 use crate::prelude::*;
 use crate::FormatNodeRule;
 use ruff_formatter::{write, FormatContext};
@@ -22,6 +25,12 @@ impl FormatNodeRule<ExprName> for FormatExprName {
     }
 }
 
+impl NeedsParentheses for ExprName {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use ruff_text_size::{TextRange, TextSize};

crates/ruff_python_formatter/src/expression/expr_named_expr.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::ExprNamedExpr;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprNamedExpr> for FormatExprNamedExpr {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprNamedExpr {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}

crates/ruff_python_formatter/src/expression/expr_set.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::ExprSet;
@@ -10,3 +13,12 @@ impl FormatNodeRule<ExprSet> for FormatExprSet {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprSet {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        match default_expression_needs_parentheses(self.into(), parenthesize, source) {
+            Parentheses::Optional => Parentheses::Never,
+            parentheses => parentheses,
+        }
+    }
+}

crates/ruff_python_formatter/src/expression/expr_set_comp.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::ExprSetComp;
@@ -10,3 +13,12 @@ impl FormatNodeRule<ExprSetComp> for FormatExprSetComp {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprSetComp {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        match default_expression_needs_parentheses(self.into(), parenthesize, source) {
+            Parentheses::Optional => Parentheses::Never,
+            parentheses => parentheses,
+        }
+    }
+}

crates/ruff_python_formatter/src/expression/expr_slice.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::ExprSlice;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprSlice> for FormatExprSlice {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprSlice {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}

crates/ruff_python_formatter/src/expression/expr_starred.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::ExprStarred;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprStarred> for FormatExprStarred {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprStarred {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}

crates/ruff_python_formatter/src/expression/expr_subscript.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::ExprSubscript;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprSubscript> for FormatExprSubscript {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprSubscript {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}

crates/ruff_python_formatter/src/expression/expr_tuple.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::ExprTuple;
@@ -10,3 +13,12 @@ impl FormatNodeRule<ExprTuple> for FormatExprTuple {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprTuple {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        match default_expression_needs_parentheses(self.into(), parenthesize, source) {
+            Parentheses::Optional => Parentheses::Never,
+            parentheses => parentheses,
+        }
+    }
+}

crates/ruff_python_formatter/src/expression/expr_unary_op.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::ExprUnaryOp;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprUnaryOp> for FormatExprUnaryOp {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprUnaryOp {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}

crates/ruff_python_formatter/src/expression/expr_yield.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::ExprYield;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprYield> for FormatExprYield {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprYield {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}

crates/ruff_python_formatter/src/expression/expr_yield_from.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::ExprYieldFrom;
@@ -10,3 +13,9 @@ impl FormatNodeRule<ExprYieldFrom> for FormatExprYieldFrom {
         write!(f, [verbatim_text(item.range)])
     }
 }
+
+impl NeedsParentheses for ExprYieldFrom {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        default_expression_needs_parentheses(self.into(), parenthesize, source)
+    }
+}

crates/ruff_python_formatter/src/expression/maybe_parenthesize.rs

@@ -1,53 +0,0 @@
-use crate::context::NodeLevel;
-use crate::prelude::*;
-use ruff_formatter::{format_args, write};
-use rustpython_parser::ast::Expr;
-
-/// Formats the passed expression. Adds parentheses if the expression doesn't fit on a line.
-pub(crate) const fn maybe_parenthesize(expression: &Expr) -> MaybeParenthesize {
-    MaybeParenthesize { expression }
-}
-
-pub(crate) struct MaybeParenthesize<'a> {
-    expression: &'a Expr,
-}
-
-impl Format<PyFormatContext<'_>> for MaybeParenthesize<'_> {
-    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
-        let saved_level = f.context().node_level();
-        f.context_mut().set_node_level(NodeLevel::Parenthesized);
-
-        let result = if needs_parentheses(self.expression) {
-            write!(
-                f,
-                [group(&format_args![
-                    if_group_breaks(&text("(")),
-                    soft_block_indent(&self.expression.format()),
-                    if_group_breaks(&text(")"))
-                ])]
-            )
-        } else {
-            // Don't add parentheses around expressions that  have parentheses on their own (e.g. list, dict, tuple, call expression)
-            self.expression.format().fmt(f)
-        };
-
-        f.context_mut().set_node_level(saved_level);
-
-        result
-    }
-}
-
-const fn needs_parentheses(expr: &Expr) -> bool {
-    !matches!(
-        expr,
-        Expr::Tuple(_)
-            | Expr::List(_)
-            | Expr::Set(_)
-            | Expr::Dict(_)
-            | Expr::ListComp(_)
-            | Expr::SetComp(_)
-            | Expr::DictComp(_)
-            | Expr::GeneratorExp(_)
-            | Expr::Call(_)
-    )
-}

crates/ruff_python_formatter/src/expression/mod.rs

@@ -1,5 +1,9 @@
+use crate::context::NodeLevel;
+use crate::expression::parentheses::{NeedsParentheses, Parentheses, Parenthesize};
 use crate::prelude::*;
-use ruff_formatter::{FormatOwnedWithRule, FormatRefWithRule, FormatRule};
+use ruff_formatter::{
+    format_args, write, FormatOwnedWithRule, FormatRefWithRule, FormatRule, FormatRuleWithOptions,
+};
 use rustpython_parser::ast::Expr;
 
 pub(crate) mod expr_attribute;
@@ -29,17 +33,30 @@ pub(crate) mod expr_tuple;
 pub(crate) mod expr_unary_op;
 pub(crate) mod expr_yield;
 pub(crate) mod expr_yield_from;
-pub(crate) mod maybe_parenthesize;
+pub(crate) mod parentheses;
 
 #[derive(Default)]
-pub struct FormatExpr;
+pub struct FormatExpr {
+    parenthesize: Parenthesize,
+}
+
+impl FormatRuleWithOptions<Expr, PyFormatContext<'_>> for FormatExpr {
+    type Options = Parenthesize;
+
+    fn with_options(mut self, options: Self::Options) -> Self {
+        self.parenthesize = options;
+        self
+    }
+}
 
 impl FormatRule<Expr, PyFormatContext<'_>> for FormatExpr {
     fn fmt(&self, item: &Expr, f: &mut PyFormatter) -> FormatResult<()> {
-        match item {
+        let parentheses = item.needs_parentheses(self.parenthesize, f.context().contents());
+
+        let format_expr = format_with(|f| match item {
             Expr::BoolOp(expr) => expr.format().fmt(f),
             Expr::NamedExpr(expr) => expr.format().fmt(f),
-            Expr::BinOp(expr) => expr.format().fmt(f),
+            Expr::BinOp(expr) => expr.format().with_options(Some(parentheses)).fmt(f),
             Expr::UnaryOp(expr) => expr.format().fmt(f),
             Expr::Lambda(expr) => expr.format().fmt(f),
             Expr::IfExp(expr) => expr.format().fmt(f),
@@ -64,6 +81,72 @@ impl FormatRule<Expr, PyFormatContext<'_>> for FormatExpr {
             Expr::List(expr) => expr.format().fmt(f),
             Expr::Tuple(expr) => expr.format().fmt(f),
             Expr::Slice(expr) => expr.format().fmt(f),
+        });
+
+        let saved_level = f.context().node_level();
+        f.context_mut().set_node_level(NodeLevel::Expression);
+
+        let result = match parentheses {
+            Parentheses::Always => {
+                write!(
+                    f,
+                    [group(&format_args![
+                        text("("),
+                        soft_block_indent(&format_expr),
+                        text(")")
+                    ])]
+                )
+            }
+            // Add optional parentheses. Ignore if the item renders parentheses itself.
+            Parentheses::Optional => {
+                write!(
+                    f,
+                    [group(&format_args![
+                        if_group_breaks(&text("(")),
+                        soft_block_indent(&format_expr),
+                        if_group_breaks(&text(")"))
+                    ])]
+                )
+            }
+            Parentheses::Custom | Parentheses::Never => Format::fmt(&format_expr, f),
+        };
+
+        f.context_mut().set_node_level(saved_level);
+
+        result
+    }
+}
+
+impl NeedsParentheses for Expr {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses {
+        match self {
+            Expr::BoolOp(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::NamedExpr(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::BinOp(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::UnaryOp(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::Lambda(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::IfExp(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::Dict(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::Set(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::ListComp(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::SetComp(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::DictComp(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::GeneratorExp(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::Await(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::Yield(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::YieldFrom(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::Compare(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::Call(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::FormattedValue(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::JoinedStr(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::Constant(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::Attribute(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::Subscript(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::Starred(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::Name(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::List(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::Tuple(expr) => expr.needs_parentheses(parenthesize, source),
+            Expr::Slice(expr) => expr.needs_parentheses(parenthesize, source),
         }
     }
 }

crates/ruff_python_formatter/src/expression/parentheses.rs

@@ -0,0 +1,93 @@
+use crate::trivia::{
+    find_first_non_trivia_character_after, find_first_non_trivia_character_before,
+};
+use ruff_python_ast::node::AnyNodeRef;
+
+pub(crate) trait NeedsParentheses {
+    fn needs_parentheses(&self, parenthesize: Parenthesize, source: &str) -> Parentheses;
+}
+
+pub(super) fn default_expression_needs_parentheses(
+    node: AnyNodeRef,
+    parenthesize: Parenthesize,
+    source: &str,
+) -> Parentheses {
+    debug_assert!(
+        node.is_expression(),
+        "Should only be called for expressions"
+    );
+
+    // `Optional` or `Preserve` and expression has parentheses in source code.
+    if !parenthesize.is_if_breaks() && is_expression_parenthesized(node, source) {
+        Parentheses::Always
+    }
+    // `Optional` or `IfBreaks`: Add parentheses if the expression doesn't fit on a line
+    else if !parenthesize.is_preserve() {
+        Parentheses::Optional
+    } else {
+        //`Preserve` and expression has no parentheses in the source code
+        Parentheses::Never
+    }
+}
+
+/// Configures if the expression should be parenthesized.
+#[derive(Copy, Clone, Debug, Default)]
+pub enum Parenthesize {
+    /// Parenthesize the expression if it has parenthesis in the source.
+    #[default]
+    Preserve,
+
+    /// Parenthesizes the expression if it doesn't fit on a line OR if the expression is parenthesized in the source code.
+    Optional,
+
+    /// Parenthesizes the expression only if it doesn't fit on a line.
+    IfBreaks,
+}
+
+impl Parenthesize {
+    const fn is_if_breaks(self) -> bool {
+        matches!(self, Parenthesize::IfBreaks)
+    }
+
+    const fn is_preserve(self) -> bool {
+        matches!(self, Parenthesize::Preserve)
+    }
+}
+
+/// Whether it is necessary to add parentheses around an expression.
+/// This is different from [`Parenthesize`] in that it is the resolved representation: It takes into account
+/// whether there are parentheses in the source code or not.
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub enum Parentheses {
+    /// Always create parentheses
+    Always,
+
+    /// Only add parentheses when necessary because the expression breaks over multiple lines.
+    Optional,
+
+    /// Custom handling by the node's formatter implementation
+    Custom,
+
+    /// Never add parentheses
+    Never,
+}
+
+fn is_expression_parenthesized(expr: AnyNodeRef, contents: &str) -> bool {
+    use rustpython_parser::ast::Ranged;
+
+    debug_assert!(
+        expr.is_expression(),
+        "Should only be called for expressions"
+    );
+
+    // Search backwards to avoid ambiguity with `(a, )` and because it's faster
+    matches!(
+        find_first_non_trivia_character_after(expr.end(), contents),
+        Some((_, ')'))
+    )
+        // Search forwards to confirm that this is not a nested expression `(5 + d * 3)`
+        && matches!(
+        find_first_non_trivia_character_before(expr.start(), contents),
+        Some((_, '('))
+    )
+}