Add formatting for StmtMatch

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/match.py

@@ -0,0 +1,57 @@
+# leading match comment
+match foo: # dangling match comment
+ case "bar":
+ pass
+
+
+# leading match comment
+match ( # leading expr comment
+ # another leading expr comment
+ foo # trailing expr comment
+ # another trailing expr comment
+): # dangling match comment
+ case "bar":
+ pass
+
+
+# leading match comment
+match ( # hello
+ foo # trailing expr comment
+ , # another
+): # dangling match comment
+ case "bar":
+ pass
+
+
+match [ # comment
+ first,
+ second,
+ third
+]: # another comment
+ case ["a", "b", "c"]:
+ pass
+
+match ( # comment
+ "a b c"
+).split(): # another comment
+ case ["a", "b", "c"]:
+ pass
+
+
+match ( # comment
+ # let's go
+ yield foo
+): # another comment
+ case ["a", "b", "c"]:
+ pass
+
+
+match aaaaaaaaahhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhh: # comment
+ case "sshhhhhhhh":
+ pass
+
+
+def foo():
+ match inside_func: # comment
+ case "bar":
+ pass

crates/ruff_python_formatter/src/comments/placement.rs

@@ -62,7 +62,6 @@ pub(super) fn place_comment<'a>(
  CommentPlacement::Default(comment)
  }
  }
- AnyNodeRef::MatchCase(match_case) => handle_match_comment(comment, match_case, locator),
  AnyNodeRef::ModModule(_) => {
  handle_module_level_own_line_comment_before_class_or_function_comment(comment, locator)
  }
@@ -209,6 +208,10 @@ fn is_first_statement_in_body(statement: AnyNodeRef, has_body: AnyNodeRef) -> bo
  are_same_optional(statement, body.first())
  }
 
+ AnyNodeRef::StmtMatch(ast::StmtMatch { cases, .. }) => {
+ are_same_optional(statement, cases.first())
+ }
+
  _ => false,
  }
 }
@@ -406,105 +409,6 @@ fn handle_own_line_comment_in_clause<'a>(
  CommentPlacement::Default(comment)
 }
 
-/// Handles leading comments in front of a match case or a trailing comment of the `match` statement.
-/// ```python
-/// match pt:
-/// # Leading `case(x, y)` comment
-/// case (x, y):
-/// return Point3d(x, y, 0)
-/// # Leading `case (x, y, z)` comment
-/// case _:
-/// ```
-fn handle_match_comment<'a>(
- comment: DecoratedComment<'a>,
- match_case: &'a MatchCase,
- locator: &Locator,
-) -> CommentPlacement<'a> {
- // Must be an own line comment after the last statement in a match case
- if comment.line_position().is_end_of_line() || comment.following_node().is_some() {
- return CommentPlacement::Default(comment);
- }
-
- // And its parent match statement.
- let Some(match_stmt) = comment.enclosing_parent().and_then(AnyNodeRef::stmt_match) else {
- return CommentPlacement::Default(comment);
- };
-
- // Get the next sibling (sibling traversal would be really nice)
- let current_case_index = match_stmt
- .cases
- .iter()
- .position(|case| case == match_case)
- .expect("Expected case to belong to parent match statement.");
-
- let next_case = match_stmt.cases.get(current_case_index + 1);
-
- let comment_indentation = indentation_at_offset(comment.slice().range().start(), locator)
- .unwrap_or_default()
- .len();
- let match_case_indentation = indentation(locator, match_case).unwrap().len();
-
- if let Some(next_case) = next_case {
- // The comment's indentation is less or equal to the `case` indention and there's a following
- // `case` arm.
- // ```python
- // match pt:
- // case (x, y):
- // return Point3d(x, y, 0)
- // # Leading `case (x, y, z)` comment
- // case _:
- // pass
- // ```
- // Attach the `comment` as leading comment to the next case.
- if comment_indentation <= match_case_indentation {
- CommentPlacement::leading(next_case, comment)
- } else {
- // Otherwise, delegate to `handle_trailing_body_comment`
- // ```python
- // match pt:
- // case (x, y):
- // return Point3d(x, y, 0)
- // # Trailing case body comment
- // case _:
- // pass
- // ```
- CommentPlacement::Default(comment)
- }
- } else {
- // Comment after the last statement in a match case...
- let match_stmt_indentation = indentation(locator, match_stmt).unwrap_or_default().len();
-
- if comment_indentation <= match_case_indentation
- && comment_indentation > match_stmt_indentation
- {
- // The comment's indent matches the `case` indent (or is larger than the `match`'s indent).
- // ```python
- // match pt:
- // case (x, y):
- // return Point3d(x, y, 0)
- // case _:
- // pass
- // # Trailing match comment
- // ```
- // This is a trailing comment of the last case.
- CommentPlacement::trailing(match_case, comment)
- } else {
- // Delegate to `handle_trailing_body_comment` because it's either a trailing indent
- // for the last statement in the `case` body or a comment for the parent of the `match`
- //
- // ```python
- // match pt:
- // case (x, y):
- // return Point3d(x, y, 0)
- // case _:
- // pass
- // # trailing case comment
- // ```
- CommentPlacement::Default(comment)
- }
- }
-}
-
 /// Determine where to attach an own line comment after a branch depending on its indentation
 fn handle_own_line_comment_after_branch<'a>(
  comment: DecoratedComment<'a>,

crates/ruff_python_formatter/src/comments/visitor.rs

@@ -73,7 +73,6 @@ impl<'a> CommentsVisitor<'a> {
  enclosing: enclosing_node,
  preceding: self.preceding_node,
  following: Some(node),
- parent: self.parents.iter().rev().nth(1).copied(),
  line_position: text_position(*comment_range, self.source_code),
  slice: self.source_code.slice(*comment_range),
  };
@@ -131,7 +130,6 @@ impl<'a> CommentsVisitor<'a> {
  let comment = DecoratedComment {
  enclosing: node,
  preceding: self.preceding_node,
- parent: self.parents.last().copied(),
  following: None,
  line_position: text_position(*comment_range, self.source_code),
  slice: self.source_code.slice(*comment_range),
@@ -340,7 +338,6 @@ pub(super) struct DecoratedComment<'a> {
  enclosing: AnyNodeRef<'a>,
  preceding: Option<AnyNodeRef<'a>>,
  following: Option<AnyNodeRef<'a>>,
- parent: Option<AnyNodeRef<'a>>,
  line_position: CommentLinePosition,
  slice: SourceCodeSlice,
 }
@@ -366,11 +363,6 @@ impl<'a> DecoratedComment<'a> {
  self.enclosing
  }
 
- /// Returns the parent of the enclosing node, if any
- pub(super) fn enclosing_parent(&self) -> Option<AnyNodeRef<'a>> {
- self.parent
- }
-
  /// Returns the slice into the source code.
  pub(super) fn slice(&self) -> &SourceCodeSlice {
  &self.slice

crates/ruff_python_formatter/src/other/match_case.rs

@@ -1,12 +1,45 @@
-use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
 use ruff_formatter::{write, Buffer, FormatResult};
 use ruff_python_ast::MatchCase;
 
+use crate::expression::maybe_parenthesize_expression;
+use crate::expression::parentheses::Parenthesize;
+use crate::not_yet_implemented_custom_text;
+use crate::prelude::*;
+use crate::{FormatNodeRule, PyFormatter};
+
 #[derive(Default)]
 pub struct FormatMatchCase;
 
 impl FormatNodeRule<MatchCase> for FormatMatchCase {
  fn fmt_fields(&self, item: &MatchCase, f: &mut PyFormatter) -> FormatResult<()> {
- write!(f, [not_yet_implemented(item)])
+ let MatchCase {
+ range: _,
+ pattern: _,
+ guard,
+ body,
+ } = item;
+
+ write!(
+ f,
+ [
+ text("case"),
+ space(),
+ not_yet_implemented_custom_text("NOT_YET_IMPLEMENTED_Pattern"),
+ ]
+ )?;
+
+ if let Some(guard) = guard {
+ write!(
+ f,
+ [
+ space(),
+ text("if"),
+ space(),
+ maybe_parenthesize_expression(guard, item, Parenthesize::IfBreaks)
+ ]
+ )?;
+ }
+
+ write!(f, [text(":"), block_indent(&body.format())])
  }
 }

crates/ruff_python_formatter/src/statement/stmt_match.rs

@@ -1,12 +1,49 @@
-use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
 use ruff_formatter::{write, Buffer, FormatResult};
 use ruff_python_ast::StmtMatch;
 
+use crate::comments::trailing_comments;
+use crate::expression::maybe_parenthesize_expression;
+use crate::expression::parentheses::Parenthesize;
+use crate::prelude::*;
+use crate::{FormatNodeRule, PyFormatter};
+
 #[derive(Default)]
 pub struct FormatStmtMatch;
 
 impl FormatNodeRule<StmtMatch> for FormatStmtMatch {
  fn fmt_fields(&self, item: &StmtMatch, f: &mut PyFormatter) -> FormatResult<()> {
- write!(f, [not_yet_implemented(item)])
+ let StmtMatch {
+ range: _,
+ subject,
+ cases,
+ } = item;
+
+ let comments = f.context().comments().clone();
+ let dangling_item_comments = comments.dangling_comments(item);
+
+ // There can be at most one dangling comment after the colon in a match statement.
+ debug_assert!(dangling_item_comments.len() <= 1);
+
+ write!(
+ f,
+ [
+ text("match"),
+ space(),
+ maybe_parenthesize_expression(subject, item, Parenthesize::IfBreaks),
+ text(":"),
+ trailing_comments(dangling_item_comments)
+ ]
+ )?;
+
+ for case in cases {
+ write!(f, [block_indent(&case.format())])?;
+ }
+
+ Ok(())
+ }
+
+ fn fmt_dangling_comments(&self, _node: &StmtMatch, _f: &mut PyFormatter) -> FormatResult<()> {
+ // Handled as part of `fmt_fields`
+ Ok(())
  }
 }