Add support for PatternMatchMapping formatting (#6836)

<!--
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 https://github.com/astral-sh/ruff/issues/6644.

## Test Plan

`cargo test`

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

@@ -344,3 +344,58 @@ match 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

crates/ruff_python_formatter/src/comments/placement.rs

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

crates/ruff_python_formatter/src/pattern/pattern_match_mapping.rs

@@ -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!(
+        let PatternMatchMapping {
-            f,
+            keys,
-            [not_yet_implemented_custom_text(
+            patterns,
-                "{\"NOT_YET_IMPLEMENTED_PatternMatchMapping\": _, 2: _}",
+            rest,
-                item
+            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(())
     }
 }
 
@@ -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()))
+}

crates/ruff_python_formatter/tests/snapshots/black_compatibility@py_310__pattern_matching_complex.py.snap

@@ -165,7 +165,7 @@ match x:
          y = 0
  # case black_test_patma_073
  match x:
-@@ -16,23 +16,23 @@
+@@ -16,11 +16,11 @@
          y = 1
  # case black_test_patma_006
  match 3:
@@ -179,15 +179,9 @@ match x:
          y = 0
  # case black_check_sequence_then_mapping
  match x:
-     case [*_]:
+@@ -32,7 +32,7 @@
-         return "seq"
--    case {}:
-+    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
-         return "map"
- # case black_test_patma_035
  match x:
--    case {0: [1, 2, {}]}:
+     case {0: [1, 2, {}]}:
-+    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
          y = 0
 -    case {0: [1, 2, {}] | True} | {1: [[]]} | {0: [1, 2, {}]} | [] | "X" | {}:
 +    case NOT_YET_IMPLEMENTED_PatternMatchOf | (y):
@@ -203,30 +197,6 @@ match x:
          x = True
  # case black_test_patma_154
  match x:
-@@ -54,11 +54,11 @@
-         y = 0
- # case black_test_patma_134
- match x:
--    case {1: 0}:
-+    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
-         y = 0
--    case {0: 0}:
-+    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
-         y = 1
--    case {**z}:
-+    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
-         y = 2
- # case black_test_patma_185
- match Seq():
-@@ -72,7 +72,7 @@
-         y = 1
- # case black_test_patma_248
- match x:
--    case {"foo": bar}:
-+    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
-         y = bar
- # case black_test_patma_019
- match (0, 1, 2):
 @@ -132,13 +132,13 @@
          z = 0
  # case black_test_patma_042
@@ -236,8 +206,7 @@ match x:
          y = 0
  # case black_test_patma_034
  match x:
--    case {0: [1, 2, {}]}:
+     case {0: [1, 2, {}]}:
-+    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
          y = 0
 -    case {0: [1, 2, {}] | False} | {1: [[]]} | {0: [1, 2, {}]} | [] | "X" | {}:
 +    case NOT_YET_IMPLEMENTED_PatternMatchOf | (y):
@@ -277,11 +246,11 @@ match x:
 match x:
     case [*_]:
         return "seq"
-    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
+    case {}:
         return "map"
 # case black_test_patma_035
 match x:
-    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
+    case {0: [1, 2, {}]}:
         y = 0
     case NOT_YET_IMPLEMENTED_PatternMatchOf | (y):
         y = 1
@@ -305,11 +274,11 @@ match x:
         y = 0
 # case black_test_patma_134
 match x:
-    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
+    case {1: 0}:
         y = 0
-    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
+    case {0: 0}:
         y = 1
-    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
+    case {**z}:
         y = 2
 # case black_test_patma_185
 match Seq():
@@ -323,7 +292,7 @@ match x:
         y = 1
 # case black_test_patma_248
 match x:
-    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
+    case {"foo": bar}:
         y = bar
 # case black_test_patma_019
 match (0, 1, 2):
@@ -387,7 +356,7 @@ match x:
         y = 0
 # case black_test_patma_034
 match x:
-    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
+    case {0: [1, 2, {}]}:
         y = 0
     case NOT_YET_IMPLEMENTED_PatternMatchOf | (y):
         y = 1

crates/ruff_python_formatter/tests/snapshots/black_compatibility@py_310__pattern_matching_extras.py.snap

@@ -188,15 +188,6 @@ match bar1:
          pass
      case _:
          pass
-@@ -48,7 +57,7 @@
- match a, *b, c:
-     case [*_]:
-         assert "seq" == _
--    case {}:
-+    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
-         assert "map" == b
- 
- 
 @@ -59,12 +68,7 @@
      ),
      case,
@@ -211,22 +202,20 @@ match bar1:
          pass
  
      case [a as match]:
-@@ -85,12 +89,9 @@
+@@ -87,10 +91,10 @@
- 
- 
  match something:
--    case {
+     case {
--        "key": key as key_1,
+         "key": key as key_1,
 -        "password": PASS.ONE | PASS.TWO | PASS.THREE as password,
--    }:
++        "password": NOT_YET_IMPLEMENTED_PatternMatchOf | (y) as password,
-+    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
+     }:
          pass
 -    case {"maybe": something(complicated as this) as that}:
-+    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
++    case {"maybe": NOT_YET_IMPLEMENTED_PatternMatchClass(0, 0) as that}:
          pass
  
  
-@@ -101,19 +102,17 @@
+@@ -101,19 +105,17 @@
      case 2 as b, 3 as c:
          pass
  
@@ -313,7 +302,7 @@ match (
 match a, *b, c:
     case [*_]:
         assert "seq" == _
-    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
+    case {}:
         assert "map" == b
 
 
@@ -345,9 +334,12 @@ match a, *b(), c:
 
 
 match something:
-    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
+    case {
+        "key": key as key_1,
+        "password": NOT_YET_IMPLEMENTED_PatternMatchOf | (y) as password,
+    }:
         pass
-    case {"NOT_YET_IMPLEMENTED_PatternMatchMapping": _, 2: _}:
+    case {"maybe": NOT_YET_IMPLEMENTED_PatternMatchClass(0, 0) as that}:
         pass
 
 

crates/ruff_python_formatter/tests/snapshots/format@statement__match.py.snap

@@ -350,6 +350,61 @@ match 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
 ```
 
 ## Output
@@ -719,6 +774,57 @@ match 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: value  # key  # colon  # 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
 ```