use ruff_formatter::{FormatOwnedWithRule, FormatRefWithRule, FormatRule, FormatRuleWithOptions}; use ruff_python_ast::node::AnyNodeRef; use ruff_python_ast::{Pattern, Ranged}; use ruff_python_trivia::{first_non_trivia_token, SimpleToken, SimpleTokenKind, SimpleTokenizer}; use crate::expression::parentheses::{ parenthesized, NeedsParentheses, OptionalParentheses, Parentheses, }; use crate::prelude::*; pub(crate) mod pattern_match_as; pub(crate) mod pattern_match_class; pub(crate) mod pattern_match_mapping; pub(crate) mod pattern_match_or; pub(crate) mod pattern_match_sequence; pub(crate) mod pattern_match_singleton; pub(crate) mod pattern_match_star; pub(crate) mod pattern_match_value; #[derive(Copy, Clone, PartialEq, Eq, Default)] pub struct FormatPattern { parentheses: Parentheses, } impl FormatRuleWithOptions> for FormatPattern { type Options = Parentheses; fn with_options(mut self, options: Self::Options) -> Self { self.parentheses = options; self } } impl FormatRule> for FormatPattern { fn fmt(&self, pattern: &Pattern, f: &mut PyFormatter) -> FormatResult<()> { let format_pattern = format_with(|f| match pattern { Pattern::MatchValue(pattern) => pattern.format().fmt(f), Pattern::MatchSingleton(pattern) => pattern.format().fmt(f), Pattern::MatchSequence(pattern) => pattern.format().fmt(f), Pattern::MatchMapping(pattern) => pattern.format().fmt(f), Pattern::MatchClass(pattern) => pattern.format().fmt(f), Pattern::MatchStar(pattern) => pattern.format().fmt(f), Pattern::MatchAs(pattern) => pattern.format().fmt(f), Pattern::MatchOr(pattern) => pattern.format().fmt(f), }); let parenthesize = match self.parentheses { Parentheses::Preserve => is_pattern_parenthesized(pattern, f.context().source()), Parentheses::Always => true, Parentheses::Never => false, }; if parenthesize { let comments = f.context().comments().clone(); // Any comments on the open parenthesis. // // For example, `# comment` in: // ```python // ( # comment // 1 // ) // ``` let open_parenthesis_comment = comments .leading(pattern) .first() .filter(|comment| comment.line_position().is_end_of_line()); parenthesized("(", &format_pattern, ")") .with_dangling_comments( open_parenthesis_comment .map(std::slice::from_ref) .unwrap_or_default(), ) .fmt(f) } else { format_pattern.fmt(f) } } } impl<'ast> AsFormat> for Pattern { type Format<'a> = FormatRefWithRule<'a, Pattern, FormatPattern, PyFormatContext<'ast>>; fn format(&self) -> Self::Format<'_> { FormatRefWithRule::new(self, FormatPattern::default()) } } impl<'ast> IntoFormat> for Pattern { type Format = FormatOwnedWithRule>; fn into_format(self) -> Self::Format { FormatOwnedWithRule::new(self, FormatPattern::default()) } } fn is_pattern_parenthesized(pattern: &Pattern, contents: &str) -> bool { // First test if there's a closing parentheses because it tends to be cheaper. if matches!( first_non_trivia_token(pattern.end(), contents), Some(SimpleToken { kind: SimpleTokenKind::RParen, .. }) ) { let mut tokenizer = SimpleTokenizer::up_to_without_back_comment(pattern.start(), contents).skip_trivia(); matches!( tokenizer.next_back(), Some(SimpleToken { kind: SimpleTokenKind::LParen, .. }) ) } else { false } } impl NeedsParentheses for Pattern { fn needs_parentheses( &self, parent: AnyNodeRef, context: &PyFormatContext, ) -> OptionalParentheses { match self { Pattern::MatchValue(pattern) => pattern.needs_parentheses(parent, context), Pattern::MatchSingleton(pattern) => pattern.needs_parentheses(parent, context), Pattern::MatchSequence(pattern) => pattern.needs_parentheses(parent, context), Pattern::MatchMapping(pattern) => pattern.needs_parentheses(parent, context), Pattern::MatchClass(pattern) => pattern.needs_parentheses(parent, context), Pattern::MatchStar(pattern) => pattern.needs_parentheses(parent, context), Pattern::MatchAs(pattern) => pattern.needs_parentheses(parent, context), Pattern::MatchOr(pattern) => pattern.needs_parentheses(parent, context), } } }