Add formatting for StmtMatch (#6286)

## Summary

This PR adds support for `StmtMatch` with subs for `MatchCase`.

## Test Plan

Add a few additional test cases around `match` statement, comments, line
breaks.

resolves: #6298
This commit is contained in:
Dhruv Manilawala 2023-08-08 18:48:49 +05:30 committed by GitHub
parent 87984e9ac7
commit 001aa486df
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
12 changed files with 882 additions and 444 deletions

View file

@ -63,7 +63,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)
}
@ -210,6 +209,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,
}
}
@ -407,105 +410,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>,

View file

@ -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

View file

@ -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())])
}
}

View file

@ -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(())
}
}