basic formatting for ExprDict (#5167)

crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/dict.py

@@ -0,0 +1,50 @@
+# before
+{  # open
+ key# key
+ : # colon
+ value# value
+} # close
+# after
+
+
+{**d}
+
+{**a, #  leading
+** # middle
+b # trailing
+}
+
+{
+** # middle with single item
+b
+}
+
+{
+    # before
+    ** # between
+    b,
+}
+
+{
+    **a # comment before preceeding node's comma
+ ,
+    # before
+    ** # between
+    b,
+}
+
+{}
+
+{1:2,}
+
+{1:2,
+ 3:4,}
+
+{asdfsadfalsdkjfhalsdkjfhalskdjfhlaksjdfhlaskjdfhlaskjdfhlaksdjfh: 1, adsfadsflasdflasdfasdfasdasdf: 2}
+
+mapping = {
+    A: 0.25 * (10.0 / 12),
+    B: 0.1 * (10.0 / 12),
+    C: 0.1 * (10.0 / 12),
+    D: 0.1 * (10.0 / 12),
+}

crates/ruff_python_formatter/src/comments/placement.rs

@@ -1,6 +1,6 @@
 use crate::comments::visitor::{CommentPlacement, DecoratedComment};
 use crate::comments::CommentTextPosition;
-use crate::trivia::{SimpleTokenizer, TokenKind};
+use crate::trivia::{SimpleTokenizer, Token, TokenKind};
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::source_code::Locator;
 use ruff_python_ast::whitespace;
@@ -29,6 +29,7 @@ pub(super) fn place_comment<'a>(
     .or_else(|comment| handle_positional_only_arguments_separator_comment(comment, locator))
     .or_else(|comment| handle_trailing_binary_expression_left_or_operator_comment(comment, locator))
     .or_else(handle_leading_function_with_decorators_comment)
+    .or_else(|comment| handle_dict_unpacking_comment(comment, locator))
 }
 
 /// Handles leading comments in front of a match case or a trailing comment of the `match` statement.
@@ -885,6 +886,76 @@ fn handle_leading_function_with_decorators_comment(comment: DecoratedComment) ->
     }
 }
 
+/// Handles comments between `**` and the variable name in dict unpacking
+/// It attaches these to the appropriate value node
+///
+/// ```python
+/// {
+///     **  # comment between `**` and the variable name
+///     value
+///     ...
+/// }
+/// ```
+fn handle_dict_unpacking_comment<'a>(
+    comment: DecoratedComment<'a>,
+    locator: &Locator,
+) -> CommentPlacement<'a> {
+    match comment.enclosing_node() {
+        // TODO: can maybe also add AnyNodeRef::Arguments here, but tricky to test due to
+        // https://github.com/astral-sh/ruff/issues/5176
+        AnyNodeRef::ExprDict(_) => {}
+        _ => {
+            return CommentPlacement::Default(comment);
+        }
+    };
+
+    // no node after our comment so we can't be between `**` and the name (node)
+    let Some(following) = comment.following_node() else {
+        return CommentPlacement::Default(comment);
+    };
+
+    // we look at tokens between the previous node (or the start of the dict)
+    // and the comment
+    let preceding_end = match comment.preceding_node() {
+        Some(preceding) => preceding.end(),
+        None => comment.enclosing_node().start(),
+    };
+    if preceding_end > comment.slice().start() {
+        return CommentPlacement::Default(comment);
+    }
+    let mut tokens = SimpleTokenizer::new(
+        locator.contents(),
+        TextRange::new(preceding_end, comment.slice().start()),
+    )
+    .skip_trivia();
+
+    // we start from the preceding node but we skip its token
+    if let Some(first) = tokens.next() {
+        debug_assert!(matches!(
+            first,
+            Token {
+                kind: TokenKind::LBrace | TokenKind::Comma | TokenKind::Colon,
+                ..
+            }
+        ));
+    }
+
+    // if the remaining tokens from the previous node is exactly `**`,
+    // re-assign the comment to the one that follows the stars
+    let mut count = 0;
+    for token in tokens {
+        if token.kind != TokenKind::Star {
+            return CommentPlacement::Default(comment);
+        }
+        count += 1;
+    }
+    if count == 2 {
+        return CommentPlacement::trailing(following, comment);
+    }
+
+    CommentPlacement::Default(comment)
+}
+
 /// Returns `true` if `right` is `Some` and `left` and `right` are referentially equal.
 fn are_same_optional<'a, T>(left: AnyNodeRef, right: Option<T>) -> bool
 where

crates/ruff_python_formatter/src/expression/expr_dict.rs

@@ -1,21 +1,108 @@
-use crate::comments::Comments;
+use crate::comments::{dangling_node_comments, leading_comments, Comments};
+use crate::context::PyFormatContext;
 use crate::expression::parentheses::{
     default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
 };
-use crate::{not_yet_implemented_custom_text, FormatNodeRule, PyFormatter};
+use crate::prelude::*;
+use crate::trivia::Token;
+use crate::trivia::{first_non_trivia_token, TokenKind};
+use crate::USE_MAGIC_TRAILING_COMMA;
+use crate::{FormatNodeRule, PyFormatter};
+use ruff_formatter::format_args;
 use ruff_formatter::{write, Buffer, FormatResult};
-use rustpython_parser::ast::ExprDict;
+use ruff_python_ast::prelude::Ranged;
+use rustpython_parser::ast::{Expr, ExprDict};
 
 #[derive(Default)]
 pub struct FormatExprDict;
 
+struct KeyValuePair<'a> {
+    key: &'a Option<Expr>,
+    value: &'a Expr,
+}
+
+impl Format<PyFormatContext<'_>> for KeyValuePair<'_> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+        if let Some(key) = self.key {
+            write!(
+                f,
+                [group(&format_args![
+                    key.format(),
+                    text(":"),
+                    space(),
+                    self.value.format()
+                ])]
+            )
+        } else {
+            let comments = f.context().comments().clone();
+            let leading_value_comments = comments.leading_comments(self.value.into());
+            write!(
+                f,
+                [
+                    // make sure the leading comments are hoisted past the `**`
+                    leading_comments(leading_value_comments),
+                    group(&format_args![text("**"), self.value.format()])
+                ]
+            )
+        }
+    }
+}
+
 impl FormatNodeRule<ExprDict> for FormatExprDict {
-    fn fmt_fields(&self, _item: &ExprDict, f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_fields(&self, item: &ExprDict, f: &mut PyFormatter) -> FormatResult<()> {
+        let ExprDict {
+            range: _,
+            keys,
+            values,
+        } = item;
+
+        let last = match &values[..] {
+            [] => {
+                return write!(
+                    f,
+                    [
+                        &text("{"),
+                        block_indent(&dangling_node_comments(item)),
+                        &text("}"),
+                    ]
+                );
+            }
+            [.., last] => last,
+        };
+        let magic_trailing_comma = USE_MAGIC_TRAILING_COMMA
+            && matches!(
+                first_non_trivia_token(last.range().end(), f.context().contents()),
+                Some(Token {
+                    kind: TokenKind::Comma,
+                    ..
+                })
+            );
+
+        debug_assert_eq!(keys.len(), values.len());
+
+        let joined = format_with(|f| {
+            f.join_with(format_args!(text(","), soft_line_break_or_space()))
+                .entries(
+                    keys.iter()
+                        .zip(values)
+                        .map(|(key, value)| KeyValuePair { key, value }),
+                )
+                .finish()
+        });
+
+        let block = if magic_trailing_comma {
+            block_indent
+        } else {
+            soft_block_indent
+        };
+
         write!(
             f,
-            [not_yet_implemented_custom_text(
-                "{NOT_IMPLEMENTED_dict_key: NOT_IMPLEMENTED_dict_value}"
-            )]
+            [group(&format_args![
+                text("{"),
+                block(&format_args![joined, if_group_breaks(&text(",")),]),
+                text("}")
+            ])]
         )
     }
 

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__collections_py.snap

@@ -84,7 +84,7 @@ if True:
 ```diff
 --- Black
 +++ Ruff
-@@ -1,99 +1,52 @@
+@@ -1,99 +1,56 @@
 -import core, time, a
 +NOT_YET_IMPLEMENTED_StmtImport
  
@@ -143,24 +143,24 @@ if True:
 -    "dddddddddddddddddddddddddddddddddddddddd",
 +    "NOT_YET_IMPLEMENTED_STRING",
  ]
--{
+ {
 -    "oneple": (1,),
--}
++    "NOT_YET_IMPLEMENTED_STRING": (1,),
+ }
 -{"oneple": (1,)}
 -["ls", "lsoneple/%s" % (foo,)]
 -x = {"oneple": (1,)}
--y = {
++{"NOT_YET_IMPLEMENTED_STRING": (1,)}
++["NOT_YET_IMPLEMENTED_STRING", "NOT_YET_IMPLEMENTED_STRING" % (foo,)]
++x = {"NOT_YET_IMPLEMENTED_STRING": (1,)}
+ y = {
 -    "oneple": (1,),
--}
++    "NOT_YET_IMPLEMENTED_STRING": (1,),
+ }
 -assert False, (
 -    "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa wraps %s"
 -    % bar
 -)
-+{NOT_IMPLEMENTED_dict_key: NOT_IMPLEMENTED_dict_value}
-+{NOT_IMPLEMENTED_dict_key: NOT_IMPLEMENTED_dict_value}
-+["NOT_YET_IMPLEMENTED_STRING", "NOT_YET_IMPLEMENTED_STRING" % (foo,)]
-+x = {NOT_IMPLEMENTED_dict_key: NOT_IMPLEMENTED_dict_value}
-+y = {NOT_IMPLEMENTED_dict_key: NOT_IMPLEMENTED_dict_value}
 +NOT_YET_IMPLEMENTED_StmtAssert
  
  # looping over a 1-tuple should also not get wrapped
@@ -245,11 +245,15 @@ nested_long_lines = [
     (1, 2, 3),
     "NOT_YET_IMPLEMENTED_STRING",
 ]
-{NOT_IMPLEMENTED_dict_key: NOT_IMPLEMENTED_dict_value}
-{NOT_IMPLEMENTED_dict_key: NOT_IMPLEMENTED_dict_value}
+{
+    "NOT_YET_IMPLEMENTED_STRING": (1,),
+}
+{"NOT_YET_IMPLEMENTED_STRING": (1,)}
 ["NOT_YET_IMPLEMENTED_STRING", "NOT_YET_IMPLEMENTED_STRING" % (foo,)]
-x = {NOT_IMPLEMENTED_dict_key: NOT_IMPLEMENTED_dict_value}
-y = {NOT_IMPLEMENTED_dict_key: NOT_IMPLEMENTED_dict_value}
+x = {"NOT_YET_IMPLEMENTED_STRING": (1,)}
+y = {
+    "NOT_YET_IMPLEMENTED_STRING": (1,),
+}
 NOT_YET_IMPLEMENTED_StmtAssert
 
 # looping over a 1-tuple should also not get wrapped