Add OptionalParentheses::MultilineIfFits

crates/ruff_formatter/src/builders.rs

@@ -2525,6 +2525,7 @@ impl<'a, Context> BestFitting<'a, Context> {
     /// # Ok(())
     /// # }
     /// ```
+    #[must_use]
     pub fn with_mode(mut self, mode: BestFittingMode) -> Self {
         self.mode = mode;
         self

crates/ruff_formatter/src/printer/mod.rs

@@ -144,12 +144,7 @@ impl<'a> Printer<'a> {
             }
 
             FormatElement::Tag(StartGroup(group)) => {
-                let print_mode =
-                    self.print_group(TagKind::Group, group.mode(), args, queue, stack)?;
-
-                if let Some(id) = group.id() {
-                    self.state.group_modes.insert_print_mode(id, print_mode);
-                }
+                self.print_group(TagKind::Group, group.mode(), group.id(), args, queue, stack)?;
             }
 
             FormatElement::Tag(StartConditionalGroup(group)) => {
@@ -160,7 +155,14 @@ impl<'a> Printer<'a> {
                 };
 
                 if expected_mode == condition.mode {
-                    self.print_group(TagKind::ConditionalGroup, group.mode(), args, queue, stack)?;
+                    self.print_group(
+                        TagKind::ConditionalGroup,
+                        group.mode(),
+                        None,
+                        args,
+                        queue,
+                        stack,
+                    )?;
                 } else {
                     // Condition isn't met, render as normal content
                     stack.push(TagKind::ConditionalGroup, args);
@@ -293,6 +295,7 @@ impl<'a> Printer<'a> {
         &mut self,
         kind: TagKind,
         mode: GroupMode,
+        id: Option<GroupId>,
         args: PrintElementArgs,
         queue: &PrintQueue<'a>,
         stack: &mut PrintCallStack,
@@ -311,6 +314,14 @@ impl<'a> Printer<'a> {
                     _ => {
                         self.state.measured_group_fits = true;
 
+                        // Set the print mode for this group in case this group was measured as part of a
+                        // previous `fits_group` call.
+                        if let Some(id) = id {
+                            self.state
+                                .group_modes
+                                .insert_print_mode(id, PrintMode::Flat);
+                        }
+
                         // Measure to see if the group fits up on a single line. If that's the case,
                         // print the group in "flat" mode, otherwise continue in expanded mode
                         stack.push(kind, args.with_print_mode(PrintMode::Flat));
@@ -329,6 +340,10 @@ impl<'a> Printer<'a> {
 
         stack.push(kind, args.with_print_mode(group_mode));
 
+        if let Some(id) = id {
+            self.state.group_modes.insert_print_mode(id, group_mode);
+        }
+
         Ok(group_mode)
     }
 
@@ -785,15 +800,14 @@ impl GroupModes {
         self.0[index] = Some(mode);
     }
 
-    fn get_print_mode(&self, group_id: GroupId) -> Option<PrintMode> {
+    fn unwrap_print_mode(&self, group_id: GroupId, next_element: &FormatElement) -> PrintMode {
         let index = u32::from(group_id) as usize;
-        self.0
+        let print_mode = self
+            .0
             .get(index)
-            .and_then(|option| option.as_ref().copied())
-    }
+            .and_then(|option| option.as_ref().copied());
 
-    fn unwrap_print_mode(&self, group_id: GroupId, next_element: &FormatElement) -> PrintMode {
-        self.get_print_mode(group_id).unwrap_or_else(|| {
+        print_mode.unwrap_or_else(|| {
             panic!("Expected group with id {group_id:?} to exist but it wasn't present in the document. Ensure that a group with such a document appears in the document before the element {next_element:?}.")
         })
     }
@@ -1123,10 +1137,7 @@ impl<'a, 'print> FitsMeasurer<'a, 'print> {
 
                 let print_mode = match condition.group_id {
                     None => args.mode(),
-                    Some(group_id) => self
-                        .group_modes()
-                        .get_print_mode(group_id)
-                        .unwrap_or_else(|| args.mode()),
+                    Some(group_id) => self.group_modes().unwrap_print_mode(group_id, element),
                 };
 
                 if condition.mode == print_mode {
@@ -1143,10 +1154,7 @@ impl<'a, 'print> FitsMeasurer<'a, 'print> {
             FormatElement::Tag(StartConditionalContent(condition)) => {
                 let print_mode = match condition.group_id {
                     None => args.mode(),
-                    Some(group_id) => self
-                        .group_modes()
-                        .get_print_mode(group_id)
-                        .unwrap_or_else(|| args.mode()),
+                    Some(group_id) => self.group_modes().unwrap_print_mode(group_id, element),
                 };
 
                 if condition.mode == print_mode {
@@ -1157,10 +1165,7 @@ impl<'a, 'print> FitsMeasurer<'a, 'print> {
             }
 
             FormatElement::Tag(StartIndentIfGroupBreaks(id)) => {
-                let print_mode = self
-                    .group_modes()
-                    .get_print_mode(*id)
-                    .unwrap_or_else(|| args.mode());
+                let print_mode = self.group_modes().unwrap_print_mode(*id, element);
 
                 match print_mode {
                     PrintMode::Flat => {
@@ -1191,10 +1196,9 @@ impl<'a, 'print> FitsMeasurer<'a, 'print> {
                 let condition_met = match condition {
                     Some(condition) => {
                         let group_mode = match condition.group_id {
-                            Some(group_id) => self
-                                .group_modes()
-                                .get_print_mode(group_id)
-                                .unwrap_or_else(|| args.mode()),
+                            Some(group_id) => {
+                                self.group_modes().unwrap_print_mode(group_id, element)
+                            }
                             None => args.mode(),
                         };
 
@@ -1250,17 +1254,17 @@ impl<'a, 'print> FitsMeasurer<'a, 'print> {
     fn fits_group(
         &mut self,
         kind: TagKind,
-        mode: GroupMode,
+        group_mode: GroupMode,
         id: Option<GroupId>,
         args: PrintElementArgs,
     ) -> Fits {
-        if self.must_be_flat && !mode.is_flat() {
+        if self.must_be_flat && !group_mode.is_flat() {
             return Fits::No;
         }
 
         // Continue printing groups in expanded mode if measuring a `best_fitting` element where
         // a group expands.
-        let print_mode = if mode.is_flat() {
+        let print_mode = if group_mode.is_flat() {
             args.mode()
         } else {
             PrintMode::Expanded

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

@@ -140,3 +140,9 @@
 # Regression: https://github.com/astral-sh/ruff/issues/5338
 if a and not aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa & aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa:
     ...
+
+if (
+  not
+  # comment
+  a):
+    ...

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

@@ -0,0 +1,54 @@
+x = aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
+
+x_aa = aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
+xxxxx = aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
+
+while (
+    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
+):
+    pass
+
+while aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa:
+    pass
+
+# Only applies in `Parenthesize::IfBreaks` positions
+raise aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
+
+raise (
+    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
+)
+
+raise a from aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
+raise a from aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
+
+# Can never apply on expression statement level
+aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
+
+# Is it only relevant for items that can't break
+
+aaaaaaa = 111111111111111111111111111111111111111111111111111111111111111111111111111111
+aaaaaaa = (
+    1111111111111111111111111111111111111111111111111111111111111111111111111111111
+)
+
+aaaaaaa = """111111111111111111111111111111111111111111111111111111111111111111111111111
+1111111111111111111111111111111111111111111111111111111111111111111111111111111111111"""
+
+# Never parenthesize multiline strings
+aaaaaaa = (
+    """1111111111111111111111111111111111111111111111111111111111111111111111111111
+1111111111111111111111111111111111111111111111111111111111111111111111111111111111111"""
+)
+
+
+
+aaaaaaaa = aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbb
+aaaaaaaa = aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
+
+aaaaaaaa = aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
+
+
+for converter in connection.ops.get_db_converters(
+    expression
+) + expression.get_db_converters(connection):
+    ...

crates/ruff_python_formatter/src/expression/expr_call.rs

@@ -69,7 +69,10 @@ impl NeedsParentheses for ExprCall {
         {
             OptionalParentheses::Multiline
         } else {
-            self.func.needs_parentheses(self.into(), context)
+            match self.func.needs_parentheses(self.into(), context) {
+                OptionalParentheses::IfFits => OptionalParentheses::Never,
+                parentheses => parentheses,
+            }
         }
     }
 }

crates/ruff_python_formatter/src/expression/expr_constant.rs

@@ -79,15 +79,12 @@ impl NeedsParentheses for ExprConstant {
         _parent: AnyNodeRef,
         context: &PyFormatContext,
     ) -> OptionalParentheses {
-        if self.value.is_implicit_concatenated() {
-            // Don't wrap triple quoted strings
-            if is_multiline_string(self, context.source()) {
-                OptionalParentheses::Never
-            } else {
-                OptionalParentheses::Multiline
-            }
-        } else {
+        if is_multiline_string(self, context.source()) {
             OptionalParentheses::Never
+        } else if self.value.is_implicit_concatenated() {
+            OptionalParentheses::Multiline
+        } else {
+            OptionalParentheses::IfFits
         }
     }
 }

crates/ruff_python_formatter/src/expression/expr_name.rs

@@ -40,7 +40,7 @@ impl NeedsParentheses for ExprName {
         _parent: AnyNodeRef,
         _context: &PyFormatContext,
     ) -> OptionalParentheses {
-        OptionalParentheses::Never
+        OptionalParentheses::IfFits
     }
 }
 

crates/ruff_python_formatter/src/expression/expr_subscript.rs

@@ -101,7 +101,10 @@ impl NeedsParentheses for ExprSubscript {
             {
                 OptionalParentheses::Multiline
             } else {
-                self.value.needs_parentheses(self.into(), context)
+                match self.value.needs_parentheses(self.into(), context) {
+                    OptionalParentheses::IfFits => OptionalParentheses::Never,
+                    parentheses => parentheses,
+                }
             }
         }
     }

crates/ruff_python_formatter/src/expression/expr_unary_op.rs

@@ -3,7 +3,7 @@ use ruff_python_ast::{ExprUnaryOp, Ranged};
 use ruff_text_size::{TextLen, TextRange};
 
 use ruff_formatter::prelude::{hard_line_break, space, text};
-use ruff_formatter::{Format, FormatContext, FormatResult};
+use ruff_formatter::{Format, FormatResult};
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
 
@@ -57,7 +57,7 @@ impl FormatNodeRule<ExprUnaryOp> for FormatExprUnaryOp {
         //  a)
         // ```
         if !leading_operand_comments.is_empty()
-            && !is_operand_parenthesized(item, f.context().source_code().as_str())
+            && !is_operand_parenthesized(item, f.context().source())
         {
             hard_line_break().fmt(f)?;
         } else if op.is_not() {

crates/ruff_python_formatter/src/expression/mod.rs

@@ -1,7 +1,7 @@
 use std::cmp::Ordering;
 
 use ruff_formatter::{
-    write, FormatOwnedWithRule, FormatRefWithRule, FormatRule, FormatRuleWithOptions,
+    format_args, write, FormatOwnedWithRule, FormatRefWithRule, FormatRule, FormatRuleWithOptions,
 };
 use ruff_python_ast as ast;
 use ruff_python_ast::node::AnyNodeRef;
@@ -223,6 +223,35 @@ impl Format<PyFormatContext<'_>> for MaybeParenthesizeExpression<'_> {
                     }
                 }
             },
+            OptionalParentheses::IfFits => match parenthesize {
+                Parenthesize::IfBreaksOrIfRequired => {
+                    parenthesize_if_expands(&expression.format().with_options(Parentheses::Never))
+                        .fmt(f)
+                }
+
+                Parenthesize::Optional | Parenthesize::IfRequired => {
+                    expression.format().with_options(Parentheses::Never).fmt(f)
+                }
+                Parenthesize::IfBreaks => {
+                    let format_expression = expression
+                        .format()
+                        .with_options(Parentheses::Never)
+                        .memoized();
+
+                    best_fitting![
+                        format_expression,
+                        group(&format_args![
+                            text("("),
+                            soft_block_indent(&format_expression),
+                            text(")")
+                        ])
+                        .should_expand(true),
+                        format_expression
+                    ]
+                    .with_mode(BestFittingMode::AllLines)
+                    .fmt(f)
+                }
+            },
             OptionalParentheses::Never => match parenthesize {
                 Parenthesize::IfBreaksOrIfRequired => {
                     parenthesize_if_expands(&expression.format().with_options(Parentheses::Never))
@@ -233,6 +262,7 @@ impl Format<PyFormatContext<'_>> for MaybeParenthesizeExpression<'_> {
                     expression.format().with_options(Parentheses::Never).fmt(f)
                 }
             },
+
             OptionalParentheses::Always => {
                 expression.format().with_options(Parentheses::Always).fmt(f)
             }

crates/ruff_python_formatter/src/expression/parentheses.rs

@@ -15,11 +15,16 @@ pub(crate) enum OptionalParentheses {
     /// Add parentheses if the expression expands over multiple lines
     Multiline,
 
-    /// Always set parentheses regardless if the expression breaks or if they were
+    /// Always set parentheses regardless if the expression breaks or if they are
     /// present in the source.
     Always,
 
-    /// Never add parentheses
+    /// Add parentheses if it helps to make this expression fit. Otherwise never add parentheses.
+    /// This mode should only be used for expressions that don't have their own split points, e.g. identifiers,
+    /// or constants.
+    IfFits,
+
+    /// Never add parentheses. Use it for expressions that have their own parentheses or if the expression body always spans multiple lines (multiline strings).
     Never,
 }
 
@@ -49,6 +54,7 @@ pub(crate) enum Parenthesize {
     /// Parenthesizes the expression if the group doesn't fit on a line (e.g., even name expressions are parenthesized), or if
     /// the expression doesn't break, but _does_ reports that it always requires parentheses in this position (e.g., walrus
     /// operators in function return annotations).
+    // TODO can this be replaced by IfBreaks
     IfBreaksOrIfRequired,
 }