Align formatting of patterns in match-cases with expression formatting in clause headers (#13510)

crates/ruff_python_formatter/src/pattern/mod.rs

@@ -1,14 +1,17 @@
 use ruff_formatter::{FormatOwnedWithRule, FormatRefWithRule, FormatRule, FormatRuleWithOptions};
-use ruff_python_ast::AnyNodeRef;
-use ruff_python_ast::Pattern;
+use ruff_python_ast::{AnyNodeRef, Expr};
+use ruff_python_ast::{MatchCase, Pattern};
 use ruff_python_trivia::CommentRanges;
 use ruff_python_trivia::{
     first_non_trivia_token, BackwardsTokenizer, SimpleToken, SimpleTokenKind,
 };
 use ruff_text_size::Ranged;
+use std::cmp::Ordering;
 
+use crate::builders::parenthesize_if_expands;
+use crate::context::{NodeLevel, WithNodeLevel};
 use crate::expression::parentheses::{
-    parenthesized, NeedsParentheses, OptionalParentheses, Parentheses,
+    optional_parentheses, parenthesized, NeedsParentheses, OptionalParentheses, Parentheses,
 };
 use crate::prelude::*;
 
@@ -150,3 +153,254 @@ impl NeedsParentheses for Pattern {
         }
     }
 }
+
+pub(crate) fn maybe_parenthesize_pattern<'a>(
+    pattern: &'a Pattern,
+    case: &'a MatchCase,
+) -> MaybeParenthesizePattern<'a> {
+    MaybeParenthesizePattern { pattern, case }
+}
+
+#[derive(Debug)]
+pub(crate) struct MaybeParenthesizePattern<'a> {
+    pattern: &'a Pattern,
+    case: &'a MatchCase,
+}
+
+impl Format<PyFormatContext<'_>> for MaybeParenthesizePattern<'_> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+        let MaybeParenthesizePattern { pattern, case } = self;
+
+        let comments = f.context().comments();
+        let pattern_comments = comments.leading_dangling_trailing(*pattern);
+
+        // If the pattern has comments, we always want to preserve the parentheses. This also
+        // ensures that we correctly handle parenthesized comments, and don't need to worry about
+        // them in the implementation below.
+        if pattern_comments.has_leading() || pattern_comments.has_trailing_own_line() {
+            return pattern.format().with_options(Parentheses::Always).fmt(f);
+        }
+
+        let needs_parentheses = pattern.needs_parentheses(AnyNodeRef::from(*case), f.context());
+
+        match needs_parentheses {
+            OptionalParentheses::Always => {
+                pattern.format().with_options(Parentheses::Always).fmt(f)
+            }
+            OptionalParentheses::Never => pattern.format().with_options(Parentheses::Never).fmt(f),
+            OptionalParentheses::Multiline => {
+                if can_pattern_omit_optional_parentheses(pattern, f.context()) {
+                    optional_parentheses(&pattern.format().with_options(Parentheses::Never)).fmt(f)
+                } else {
+                    parenthesize_if_expands(&pattern.format().with_options(Parentheses::Never))
+                        .fmt(f)
+                }
+            }
+            OptionalParentheses::BestFit => {
+                if pattern_comments.has_trailing() {
+                    pattern.format().with_options(Parentheses::Always).fmt(f)
+                } else {
+                    // The group id is necessary because the nested expressions may reference it.
+                    let group_id = f.group_id("optional_parentheses");
+                    let f = &mut WithNodeLevel::new(NodeLevel::Expression(Some(group_id)), f);
+
+                    best_fit_parenthesize(&pattern.format().with_options(Parentheses::Never))
+                        .with_group_id(Some(group_id))
+                        .fmt(f)
+                }
+            }
+        }
+    }
+}
+
+/// This function is very similar to [`can_omit_optional_parentheses`] with the only difference that it is for patterns
+/// and not expressions.
+///
+/// The base idea of the omit optional parentheses layout is to prefer using parentheses of sub-patterns
+/// when splitting the pattern over introducing new patterns. For example, prefer splitting the sequence pattern in
+/// `a | [b, c]` over splitting before the `|` operator.
+///
+/// The layout is only applied when the parenthesized pattern is the first or last item in the pattern.
+/// For example, the layout isn't used for `a | [b, c] | d` because that would look weird.
+pub(crate) fn can_pattern_omit_optional_parentheses(
+    pattern: &Pattern,
+    context: &PyFormatContext,
+) -> bool {
+    let mut visitor = CanOmitOptionalParenthesesVisitor::default();
+    visitor.visit_pattern(pattern, context);
+
+    if !visitor.any_parenthesized_expressions {
+        // Only use the more complex IR if there's a parenthesized pattern that can be split before
+        // splitting other patterns. E.g. split the sequence pattern before the string literal `"a" "b" | [a, b, c, d]`.
+        false
+    } else if visitor.max_precedence_count > 1 {
+        false
+    } else {
+        // It's complicated
+        fn has_parentheses_and_is_non_empty(pattern: &Pattern, context: &PyFormatContext) -> bool {
+            let has_own_non_empty = match pattern {
+                Pattern::MatchValue(_)
+                | Pattern::MatchSingleton(_)
+                | Pattern::MatchStar(_)
+                | Pattern::MatchAs(_)
+                | Pattern::MatchOr(_) => false,
+                Pattern::MatchSequence(sequence) => {
+                    !sequence.patterns.is_empty() || context.comments().has_dangling(pattern)
+                }
+                Pattern::MatchMapping(mapping) => {
+                    !mapping.patterns.is_empty() || context.comments().has_dangling(pattern)
+                }
+                Pattern::MatchClass(class) => !class.arguments.patterns.is_empty(),
+            };
+
+            if has_own_non_empty {
+                true
+            } else {
+                // If the pattern has no own parentheses or it is empty (e.g. ([])), check for surrounding parentheses (that should be preserved).
+                is_pattern_parenthesized(pattern, context.comments().ranges(), context.source())
+            }
+        }
+
+        visitor
+            .last
+            .is_some_and(|last| has_parentheses_and_is_non_empty(last, context))
+            || visitor
+                .first
+                .pattern()
+                .is_some_and(|first| has_parentheses_and_is_non_empty(first, context))
+    }
+}
+
+#[derive(Debug, Default)]
+struct CanOmitOptionalParenthesesVisitor<'input> {
+    max_precedence: OperatorPrecedence,
+    max_precedence_count: usize,
+    any_parenthesized_expressions: bool,
+    last: Option<&'input Pattern>,
+    first: First<'input>,
+}
+
+impl<'a> CanOmitOptionalParenthesesVisitor<'a> {
+    fn visit_pattern(&mut self, pattern: &'a Pattern, context: &PyFormatContext) {
+        match pattern {
+            Pattern::MatchSequence(_) | Pattern::MatchMapping(_) => {
+                self.any_parenthesized_expressions = true;
+            }
+
+            Pattern::MatchValue(value) => match &*value.value {
+                Expr::StringLiteral(string) => {
+                    self.update_max_precedence(OperatorPrecedence::String, string.value.len());
+                }
+                Expr::BytesLiteral(bytes) => {
+                    self.update_max_precedence(OperatorPrecedence::String, bytes.value.len());
+                }
+                // F-strings are allowed according to python's grammar but fail with a syntax error at runtime.
+                // That's why we need to support them for formatting.
+                Expr::FString(string) => {
+                    self.update_max_precedence(
+                        OperatorPrecedence::String,
+                        string.value.as_slice().len(),
+                    );
+                }
+
+                Expr::NumberLiteral(_) | Expr::Attribute(_) | Expr::UnaryOp(_) => {
+                    // require no state update other than visit_pattern does.
+                }
+
+                // `case 4+3j:` or `case 4-3j:
+                // Can not contain arbitrary expressions. Limited to complex numbers.
+                Expr::BinOp(_) => {
+                    self.update_max_precedence(OperatorPrecedence::Additive, 1);
+                }
+
+                _ => {
+                    debug_assert!(
+                        false,
+                        "Unsupported expression in pattern mach value: {:?}",
+                        value.value
+                    );
+                }
+            },
+            Pattern::MatchClass(_) => {
+                self.any_parenthesized_expressions = true;
+
+                // The pattern doesn't start with a parentheses pattern, but with the class's identifier.
+                self.first.set_if_none(First::Token);
+            }
+            Pattern::MatchStar(_) | Pattern::MatchSingleton(_) | Pattern::MatchAs(_) => {}
+            Pattern::MatchOr(or_pattern) => {
+                self.update_max_precedence(
+                    OperatorPrecedence::Or,
+                    or_pattern.patterns.len().saturating_sub(1),
+                );
+
+                for pattern in &or_pattern.patterns {
+                    self.visit_sub_pattern(pattern, context);
+                }
+            }
+        }
+    }
+
+    fn visit_sub_pattern(&mut self, pattern: &'a Pattern, context: &PyFormatContext) {
+        self.last = Some(pattern);
+
+        // Rule only applies for non-parenthesized patterns.
+        if is_pattern_parenthesized(pattern, context.comments().ranges(), context.source()) {
+            self.any_parenthesized_expressions = true;
+        } else {
+            self.visit_pattern(pattern, context);
+        }
+
+        self.first.set_if_none(First::Pattern(pattern));
+    }
+
+    fn update_max_precedence(&mut self, precedence: OperatorPrecedence, count: usize) {
+        match self.max_precedence.cmp(&precedence) {
+            Ordering::Less => {
+                self.max_precedence_count = count;
+                self.max_precedence = precedence;
+            }
+            Ordering::Equal => {
+                self.max_precedence_count += count;
+            }
+            Ordering::Greater => {}
+        }
+    }
+}
+
+#[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Default)]
+enum OperatorPrecedence {
+    #[default]
+    None,
+    Additive,
+    Or,
+    // Implicit string concatenation
+    String,
+}
+
+#[derive(Copy, Clone, Debug, Default)]
+enum First<'a> {
+    #[default]
+    None,
+
+    /// Pattern starts with a non-parentheses token. E.g. `*x`
+    Token,
+
+    Pattern(&'a Pattern),
+}
+
+impl<'a> First<'a> {
+    #[inline]
+    fn set_if_none(&mut self, first: First<'a>) {
+        if matches!(self, First::None) {
+            *self = first;
+        }
+    }
+
+    fn pattern(self) -> Option<&'a Pattern> {
+        match self {
+            First::None | First::Token => None,
+            First::Pattern(pattern) => Some(pattern),
+        }
+    }
+}

crates/ruff_python_formatter/src/pattern/pattern_match_as.rs

@@ -5,6 +5,7 @@ use ruff_python_ast::PatternMatchAs;
 use crate::comments::dangling_comments;
 use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
 use crate::prelude::*;
+use crate::preview::is_match_case_parentheses_enabled;
 
 #[derive(Default)]
 pub struct FormatPatternMatchAs;
@@ -54,8 +55,16 @@ impl NeedsParentheses for PatternMatchAs {
     fn needs_parentheses(
         &self,
         _parent: AnyNodeRef,
-        _context: &PyFormatContext,
+        context: &PyFormatContext,
     ) -> OptionalParentheses {
-        OptionalParentheses::Multiline
+        if is_match_case_parentheses_enabled(context) {
+            if self.name.is_some() {
+                OptionalParentheses::Multiline
+            } else {
+                OptionalParentheses::BestFit
+            }
+        } else {
+            OptionalParentheses::Multiline
+        }
     }
 }

crates/ruff_python_formatter/src/pattern/pattern_match_class.rs

@@ -46,7 +46,7 @@ impl NeedsParentheses for PatternMatchClass {
         // ): ...
         // ```
         if context.comments().has_dangling(self) {
-            OptionalParentheses::Multiline
+            OptionalParentheses::Always
         } else {
             OptionalParentheses::Never
         }

crates/ruff_python_formatter/src/pattern/pattern_match_or.rs

@@ -4,9 +4,11 @@ use ruff_python_ast::PatternMatchOr;
 
 use crate::comments::leading_comments;
 use crate::expression::parentheses::{
-    in_parentheses_only_soft_line_break_or_space, NeedsParentheses, OptionalParentheses,
+    in_parentheses_only_group, in_parentheses_only_soft_line_break_or_space, NeedsParentheses,
+    OptionalParentheses,
 };
 use crate::prelude::*;
+use crate::preview::is_match_case_parentheses_enabled;
 
 #[derive(Default)]
 pub struct FormatPatternMatchOr;
@@ -41,7 +43,11 @@ impl FormatNodeRule<PatternMatchOr> for FormatPatternMatchOr {
             Ok(())
         });
 
-        inner.fmt(f)
+        if is_match_case_parentheses_enabled(f.context()) {
+            in_parentheses_only_group(&inner).fmt(f)
+        } else {
+            inner.fmt(f)
+        }
     }
 }
 

crates/ruff_python_formatter/src/pattern/pattern_match_singleton.rs

@@ -3,6 +3,7 @@ use ruff_python_ast::{PatternMatchSingleton, Singleton};
 
 use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
 use crate::prelude::*;
+use crate::preview::is_match_case_parentheses_enabled;
 
 #[derive(Default)]
 pub struct FormatPatternMatchSingleton;
@@ -21,8 +22,12 @@ impl NeedsParentheses for PatternMatchSingleton {
     fn needs_parentheses(
         &self,
         _parent: AnyNodeRef,
-        _context: &PyFormatContext,
+        context: &PyFormatContext,
     ) -> OptionalParentheses {
-        OptionalParentheses::Never
+        if is_match_case_parentheses_enabled(context) {
+            OptionalParentheses::BestFit
+        } else {
+            OptionalParentheses::Never
+        }
     }
 }

crates/ruff_python_formatter/src/pattern/pattern_match_star.rs

@@ -31,6 +31,8 @@ impl NeedsParentheses for PatternMatchStar {
         _parent: AnyNodeRef,
         _context: &PyFormatContext,
     ) -> OptionalParentheses {
+        // Doesn't matter what we return here because starred patterns can never be used
+        // outside a sequence pattern.
         OptionalParentheses::Never
     }
 }

crates/ruff_python_formatter/src/pattern/pattern_match_value.rs

@@ -3,6 +3,7 @@ use ruff_python_ast::PatternMatchValue;
 
 use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses, Parentheses};
 use crate::prelude::*;
+use crate::preview::is_match_case_parentheses_enabled;
 
 #[derive(Default)]
 pub struct FormatPatternMatchValue;
@@ -17,9 +18,13 @@ impl FormatNodeRule<PatternMatchValue> for FormatPatternMatchValue {
 impl NeedsParentheses for PatternMatchValue {
     fn needs_parentheses(
         &self,
-        _parent: AnyNodeRef,
-        _context: &PyFormatContext,
+        parent: AnyNodeRef,
+        context: &PyFormatContext,
     ) -> OptionalParentheses {
-        OptionalParentheses::Never
+        if is_match_case_parentheses_enabled(context) {
+            self.value.needs_parentheses(parent, context)
+        } else {
+            OptionalParentheses::Never
+        }
     }
 }