Skip to content

Commit

Permalink
Add support for PatternMatchMapping formatting (#6836)
Browse files Browse the repository at this point in the history
<!--
Thank you for contributing to Ruff! To help us out with reviewing,
please consider the following:

- Does this pull request include a summary of the change? (See below.)
- Does this pull request include a descriptive title?
- Does this pull request include references to any relevant issues?
-->

## Summary

Adds support for `PatternMatchMapping` -- i.e., cases like:

```python
match foo:
    case {"a": 1, "b": 2, **rest}:
        pass
```

Unfortunately, this node has _three_ kinds of dangling comments:

```python
{  # "open parenthesis comment"
   key: pattern,
   **  # end-of-line "double star comment"
   # own-line "double star comment"
   rest  # end-of-line "after rest comment"
   # own-line "after rest comment"
}
```

Some of the complexity comes from the fact that in `**rest`, `rest` is
an _identifier_, not a node, so we have to handle comments _after_ it as
dangling on the enclosing node, rather than trailing on `**rest`. (We
could change the AST to use `PatternMatchAs` there, which would be more
permissive than the grammar but not totally crazy -- `PatternMatchAs` is
used elsewhere to mean "a single identifier".)

Closes #6644.

## Test Plan

`cargo test`
  • Loading branch information
charliermarsh authored Aug 25, 2023
1 parent 813d7da commit 1044d66
Show file tree
Hide file tree
Showing 6 changed files with 407 additions and 73 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -344,3 +344,58 @@ def foo():
pass
case [(*rest), (a as b)]:
pass


match foo:
case {"a": 1, "b": 2}:
pass

case {
# own line
"a": 1, # end-of-line
# own line
"b": 2,
}:
pass

case { # open
1 # key
: # colon
value # value
}:
pass

case {**d}:
pass

case {
** # middle with single item
b
}:
pass

case {
# before
** # between
b,
}:
pass

case {
1: x,
# foo
** # bop
# before
b, # boo
# baz
}:
pass

case {
1: x
# foo
,
**
b,
}:
pass
57 changes: 57 additions & 0 deletions crates/ruff_python_formatter/src/comments/placement.rs
Original file line number Diff line number Diff line change
Expand Up @@ -221,6 +221,10 @@ fn handle_enclosed_comment<'a>(
AnyNodeRef::WithItem(_) => handle_with_item_comment(comment, locator),
AnyNodeRef::PatternMatchAs(_) => handle_pattern_match_as_comment(comment, locator),
AnyNodeRef::PatternMatchStar(_) => handle_pattern_match_star_comment(comment),
AnyNodeRef::PatternMatchMapping(pattern) => {
handle_bracketed_end_of_line_comment(comment, locator)
.or_else(|comment| handle_pattern_match_mapping_comment(comment, pattern, locator))
}
AnyNodeRef::StmtFunctionDef(_) => handle_leading_function_with_decorators_comment(comment),
AnyNodeRef::StmtClassDef(class_def) => {
handle_leading_class_with_decorators_comment(comment, class_def)
Expand Down Expand Up @@ -1280,6 +1284,59 @@ fn handle_pattern_match_star_comment(comment: DecoratedComment) -> CommentPlacem
CommentPlacement::dangling(comment.enclosing_node(), comment)
}

/// Handles trailing comments after the `**` in a pattern match item. The comments can either
/// appear between the `**` and the identifier, or after the identifier (which is just an
/// identifier, not a node).
///
/// ```python
/// case {
/// ** # dangling end of line comment
/// # dangling own line comment
/// rest # dangling end of line comment
/// # dangling own line comment
/// ): ...
/// ```
fn handle_pattern_match_mapping_comment<'a>(
comment: DecoratedComment<'a>,
pattern: &'a ast::PatternMatchMapping,
locator: &Locator,
) -> CommentPlacement<'a> {
// The `**` has to come at the end, so there can't be another node after it. (The identifier,
// like `rest` above, isn't a node.)
if comment.following_node().is_some() {
return CommentPlacement::Default(comment);
};

// If there's no rest pattern, no need to do anything special.
let Some(rest) = pattern.rest.as_ref() else {
return CommentPlacement::Default(comment);
};

// If the comment falls after the `**rest` entirely, treat it as dangling on the enclosing
// node.
if comment.start() > rest.end() {
return CommentPlacement::dangling(comment.enclosing_node(), comment);
}

// Look at the tokens between the previous node (or the start of the pattern) and the comment.
let preceding_end = match comment.preceding_node() {
Some(preceding) => preceding.end(),
None => comment.enclosing_node().start(),
};
let mut tokens = SimpleTokenizer::new(
locator.contents(),
TextRange::new(preceding_end, comment.start()),
)
.skip_trivia();

// If the remaining tokens from the previous node include `**`, mark as a dangling comment.
if tokens.any(|token| token.kind == SimpleTokenKind::DoubleStar) {
CommentPlacement::dangling(comment.enclosing_node(), comment)
} else {
CommentPlacement::Default(comment)
}
}

/// Handles comments around the `:=` token in a named expression (walrus operator).
///
/// For example, here, `# 1` and `# 2` will be marked as dangling comments on the named expression,
Expand Down
175 changes: 165 additions & 10 deletions crates/ruff_python_formatter/src/pattern/pattern_match_mapping.rs
Original file line number Diff line number Diff line change
@@ -1,23 +1,103 @@
use ruff_formatter::{write, Buffer, FormatResult};
use ruff_formatter::{format_args, write};
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::PatternMatchMapping;
use ruff_python_ast::{Expr, Identifier, Pattern, Ranged};
use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
use ruff_text_size::TextRange;

use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
use crate::comments::{leading_comments, trailing_comments, SourceComment};
use crate::expression::parentheses::{
empty_parenthesized, parenthesized, NeedsParentheses, OptionalParentheses,
};
use crate::prelude::*;
use crate::{not_yet_implemented_custom_text, FormatNodeRule, PyFormatter};

#[derive(Default)]
pub struct FormatPatternMatchMapping;

impl FormatNodeRule<PatternMatchMapping> for FormatPatternMatchMapping {
fn fmt_fields(&self, item: &PatternMatchMapping, f: &mut PyFormatter) -> FormatResult<()> {
write!(
f,
[not_yet_implemented_custom_text(
"{\"NOT_YET_IMPLEMENTED_PatternMatchMapping\": _, 2: _}",
item
)]
)
let PatternMatchMapping {
keys,
patterns,
rest,
range: _,
} = item;

debug_assert_eq!(keys.len(), patterns.len());

let comments = f.context().comments().clone();
let dangling = comments.dangling(item);

if keys.is_empty() && rest.is_none() {
return empty_parenthesized("{", dangling, "}").fmt(f);
}

// This node supports three kinds of dangling comments. Most of the complexity originates
// with the rest pattern (`{**rest}`), since we can have comments around the `**`, but
// also, the `**rest` itself is not a node (it's an identifier), so comments that trail it
// are _also_ dangling.
//
// Specifically, we have these three sources of dangling comments:
// ```python
// { # "open parenthesis comment"
// key: pattern,
// ** # end-of-line "double star comment"
// # own-line "double star comment"
// rest # end-of-line "after rest comment"
// # own-line "after rest comment"
// }
// ```
let (open_parenthesis_comments, double_star_comments, after_rest_comments) =
if let Some((double_star, rest)) = find_double_star(item, f.context().source()) {
let (open_parenthesis_comments, dangling) =
dangling.split_at(dangling.partition_point(|comment| {
comment.line_position().is_end_of_line()
&& comment.start() < double_star.start()
}));
let (double_star_comments, after_rest_comments) = dangling
.split_at(dangling.partition_point(|comment| comment.start() < rest.start()));
(
open_parenthesis_comments,
double_star_comments,
after_rest_comments,
)
} else {
(dangling, [].as_slice(), [].as_slice())
};

let format_pairs = format_with(|f| {
let mut joiner = f.join_comma_separated(item.end());

for (key, pattern) in keys.iter().zip(patterns) {
let key_pattern_pair = KeyPatternPair { key, pattern };
joiner.entry(&key_pattern_pair, &key_pattern_pair);
}

if let Some(identifier) = rest {
let rest_pattern = RestPattern {
identifier,
comments: double_star_comments,
};
joiner.entry(&rest_pattern, &rest_pattern);
}

joiner.finish()?;

trailing_comments(after_rest_comments).fmt(f)
});

parenthesized("{", &format_pairs, "}")
.with_dangling_comments(open_parenthesis_comments)
.fmt(f)
}

fn fmt_dangling_comments(
&self,
_dangling_comments: &[SourceComment],
_f: &mut PyFormatter,
) -> FormatResult<()> {
// Handled by `fmt_fields`
Ok(())
}
}

Expand All @@ -30,3 +110,78 @@ impl NeedsParentheses for PatternMatchMapping {
OptionalParentheses::Never
}
}

/// A struct to format the `rest` element of a [`PatternMatchMapping`] (e.g., `{**rest}`).
#[derive(Debug)]
struct RestPattern<'a> {
identifier: &'a Identifier,
comments: &'a [SourceComment],
}

impl Ranged for RestPattern<'_> {
fn range(&self) -> TextRange {
self.identifier.range()
}
}

impl Format<PyFormatContext<'_>> for RestPattern<'_> {
fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
write!(
f,
[
leading_comments(self.comments),
text("**"),
self.identifier.format()
]
)
}
}

/// A struct to format a key-pattern pair of a [`PatternMatchMapping`] (e.g., `{key: pattern}`).
#[derive(Debug)]
struct KeyPatternPair<'a> {
key: &'a Expr,
pattern: &'a Pattern,
}

impl Ranged for KeyPatternPair<'_> {
fn range(&self) -> TextRange {
TextRange::new(self.key.start(), self.pattern.end())
}
}

impl Format<PyFormatContext<'_>> for KeyPatternPair<'_> {
fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
write!(
f,
[group(&format_args![
self.key.format(),
text(":"),
space(),
self.pattern.format()
])]
)
}
}

/// Given a [`PatternMatchMapping`], finds the range of the `**` element in the `rest` pattern,
/// if it exists.
fn find_double_star(pattern: &PatternMatchMapping, source: &str) -> Option<(TextRange, TextRange)> {
let PatternMatchMapping {
keys: _,
patterns,
rest,
range: _,
} = pattern;

// If there's no `rest` element, there's no `**`.
let Some(rest) = rest else {
return None;
};

let mut tokenizer =
SimpleTokenizer::starts_at(patterns.last().map_or(pattern.start(), Ranged::end), source);
let double_star = tokenizer.find(|token| token.kind() == SimpleTokenKind::DoubleStar)?;

Some((double_star.range(), rest.range()))
}
Loading

0 comments on commit 1044d66

Please sign in to comment.