Add JoinCommaSeparatedBuilder (#5342)

crates/ruff_python_formatter/src/builders.rs

@@ -3,7 +3,7 @@ use crate::prelude::*;
 use crate::trivia::{first_non_trivia_token, lines_after, skip_trailing_trivia, Token, TokenKind};
 use crate::USE_MAGIC_TRAILING_COMMA;
 use ruff_formatter::write;
-use ruff_text_size::{TextRange, TextSize};
+use ruff_text_size::TextSize;
 use rustpython_parser::ast::Ranged;
 
 /// Provides Python specific extensions to [`Formatter`].
@@ -16,12 +16,21 @@ pub(crate) trait PyFormatterExtensions<'ast, 'buf> {
     /// * [`NodeLevel::CompoundStatement`]: Up to one empty line
     /// * [`NodeLevel::Expression`]: No empty lines
     fn join_nodes<'fmt>(&'fmt mut self, level: NodeLevel) -> JoinNodesBuilder<'fmt, 'ast, 'buf>;
+
+    /// A builder that separates each element by a `,` and a [`soft_line_break_or_space`].
+    /// It emits a trailing `,` that is only shown if the enclosing group expands. It forces the enclosing
+    /// group to expand if the last item has a trailing `comma` and the magical comma option is enabled.
+    fn join_comma_separated<'fmt>(&'fmt mut self) -> JoinCommaSeparatedBuilder<'fmt, 'ast, 'buf>;
 }
 
 impl<'buf, 'ast> PyFormatterExtensions<'ast, 'buf> for PyFormatter<'ast, 'buf> {
     fn join_nodes<'fmt>(&'fmt mut self, level: NodeLevel) -> JoinNodesBuilder<'fmt, 'ast, 'buf> {
         JoinNodesBuilder::new(self, level)
     }
+
+    fn join_comma_separated<'fmt>(&'fmt mut self) -> JoinCommaSeparatedBuilder<'fmt, 'ast, 'buf> {
+        JoinCommaSeparatedBuilder::new(self)
+    }
 }
 
 #[must_use = "must eventually call `finish()` on the builder."]
@@ -146,15 +155,87 @@ impl<'fmt, 'ast, 'buf> JoinNodesBuilder<'fmt, 'ast, 'buf> {
     }
 }
 
-pub(crate) fn use_magic_trailing_comma(f: &mut PyFormatter, range: TextRange) -> bool {
+pub(crate) struct JoinCommaSeparatedBuilder<'fmt, 'ast, 'buf> {
-    USE_MAGIC_TRAILING_COMMA
+    result: FormatResult<()>,
-        && matches!(
+    fmt: &'fmt mut PyFormatter<'ast, 'buf>,
-            first_non_trivia_token(range.end(), f.context().contents()),
+    last_end: Option<TextSize>,
-            Some(Token {
+}
-                kind: TokenKind::Comma,
+
-                ..
+impl<'fmt, 'ast, 'buf> JoinCommaSeparatedBuilder<'fmt, 'ast, 'buf> {
-            })
+    fn new(f: &'fmt mut PyFormatter<'ast, 'buf>) -> Self {
-        )
+        Self {
+            fmt: f,
+            result: Ok(()),
+            last_end: None,
+        }
+    }
+
+    pub(crate) fn entry<T>(
+        &mut self,
+        node: &T,
+        content: &dyn Format<PyFormatContext<'ast>>,
+    ) -> &mut Self
+    where
+        T: Ranged,
+    {
+        self.result = self.result.and_then(|_| {
+            if self.last_end.is_some() {
+                write!(self.fmt, [text(","), soft_line_break_or_space()])?;
+            }
+
+            self.last_end = Some(node.end());
+
+            content.fmt(self.fmt)
+        });
+
+        self
+    }
+
+    #[allow(unused)]
+    pub(crate) fn entries<T, I, F>(&mut self, entries: I) -> &mut Self
+    where
+        T: Ranged,
+        F: Format<PyFormatContext<'ast>>,
+        I: Iterator<Item = (T, F)>,
+    {
+        for (node, content) in entries {
+            self.entry(&node, &content);
+        }
+
+        self
+    }
+
+    pub(crate) fn nodes<'a, T, I>(&mut self, entries: I) -> &mut Self
+    where
+        T: Ranged + AsFormat<PyFormatContext<'ast>> + 'a,
+        I: Iterator<Item = &'a T>,
+    {
+        for node in entries {
+            self.entry(node, &node.format());
+        }
+
+        self
+    }
+
+    pub(crate) fn finish(&mut self) -> FormatResult<()> {
+        if let Some(last_end) = self.last_end.take() {
+            if_group_breaks(&text(",")).fmt(self.fmt)?;
+
+            if USE_MAGIC_TRAILING_COMMA
+                && matches!(
+                    first_non_trivia_token(last_end, self.fmt.context().contents()),
+                    Some(Token {
+                        kind: TokenKind::Comma,
+                        ..
+                    })
+                )
+            {
+                expand_parent().fmt(self.fmt)?;
+            }
+        }
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]

crates/ruff_python_formatter/src/expression/expr_dict.rs

@@ -1,14 +1,12 @@
-use crate::builders::use_magic_trailing_comma;
 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::prelude::*;
-use crate::{FormatNodeRule, PyFormatter};
+use crate::FormatNodeRule;
-use ruff_formatter::format_args;
+use ruff_formatter::{format_args, write};
-use ruff_formatter::{write, Buffer, FormatResult};
 use ruff_python_ast::prelude::Ranged;
+use ruff_text_size::TextRange;
 use rustpython_parser::ast::{Expr, ExprDict};
 
 #[derive(Default)]
@@ -19,6 +17,16 @@ struct KeyValuePair<'a> {
     value: &'a Expr,
 }
 
+impl Ranged for KeyValuePair<'_> {
+    fn range(&self) -> TextRange {
+        if let Some(key) = self.key {
+            TextRange::new(key.start(), self.value.end())
+        } else {
+            self.value.range()
+        }
+    }
+}
+
 impl Format<PyFormatContext<'_>> for KeyValuePair<'_> {
     fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
         if let Some(key) = self.key {
@@ -54,57 +62,42 @@ impl FormatNodeRule<ExprDict> for FormatExprDict {
             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(f, last.range());
-
         debug_assert_eq!(keys.len(), values.len());
 
-        let joined = format_with(|f| {
+        if values.is_empty() {
-            f.join_with(format_args!(text(","), soft_line_break_or_space()))
+            return write!(
-                .entries(
+                f,
-                    keys.iter()
+                [
-                        .zip(values)
+                    &text("{"),
-                        .map(|(key, value)| KeyValuePair { key, value }),
+                    block_indent(&dangling_node_comments(item)),
-                )
+                    &text("}"),
-                .finish()
+                ]
-        });
+            );
+        }
 
-        let block = if magic_trailing_comma {
+        let format_pairs = format_with(|f| {
-            block_indent
+            let mut joiner = f.join_comma_separated();
-        } else {
+
-            soft_block_indent
+            for (key, value) in keys.iter().zip(values) {
-        };
+                let key_value_pair = KeyValuePair { key, value };
+                joiner.entry(&key_value_pair, &key_value_pair);
+            }
+
+            joiner.finish()
+        });
 
         write!(
             f,
             [group(&format_args![
                 text("{"),
-                block(&format_args![joined, if_group_breaks(&text(",")),]),
+                soft_block_indent(&format_pairs),
                 text("}")
             ])]
         )
     }
 
     fn fmt_dangling_comments(&self, _node: &ExprDict, _f: &mut PyFormatter) -> FormatResult<()> {
-        // TODO(konstin): Reactivate when string formatting works, currently a source of unstable
+        // Handled by `fmt_fields`
-        // formatting, e.g.
-        // ```python
-        // coverage_ignore_c_items = {
-        // #    'cfunction': [...]
-        // }
-        // ```
         Ok(())
     }
 }

crates/ruff_python_formatter/src/expression/expr_tuple.rs

@@ -1,14 +1,12 @@
-use crate::builders::use_magic_trailing_comma;
+use crate::builders::PyFormatterExtensions;
 use crate::comments::{dangling_node_comments, Comments};
 use crate::context::PyFormatContext;
 use crate::expression::parentheses::{
     default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
 };
-use crate::{AsFormat, FormatNodeRule, FormattedIterExt, PyFormatter};
+use crate::{AsFormat, FormatNodeRule, PyFormatter};
 use ruff_formatter::formatter::Formatter;
-use ruff_formatter::prelude::{
+use ruff_formatter::prelude::{block_indent, group, if_group_breaks, soft_block_indent, text};
-    block_indent, group, if_group_breaks, soft_block_indent, soft_line_break_or_space, text,
-};
 use ruff_formatter::{format_args, write, Buffer, Format, FormatResult, FormatRuleWithOptions};
 use ruff_python_ast::prelude::{Expr, Ranged};
 use ruff_text_size::TextRange;
@@ -61,9 +59,9 @@ impl FormatNodeRule<ExprTuple> for FormatExprTuple {
         } = item;
 
         // Handle the edge cases of an empty tuple and a tuple with one element
-        let last = match &elts[..] {
+        match elts.as_slice() {
             [] => {
-                return write!(
+                write!(
                     f,
                     [
                         // An empty tuple always needs parentheses, but does not have a comma
@@ -71,10 +69,10 @@ impl FormatNodeRule<ExprTuple> for FormatExprTuple {
                         block_indent(&dangling_node_comments(item)),
                         &text(")"),
                     ]
-                );
+                )
             }
             [single] => {
-                return write!(
+                write!(
                     f,
                     [group(&format_args![
                         // A single element tuple always needs parentheses and a trailing comma
@@ -82,55 +80,38 @@ impl FormatNodeRule<ExprTuple> for FormatExprTuple {
                         soft_block_indent(&format_args![single.format(), &text(",")]),
                         &text(")"),
                     ])]
-                );
+                )
             }
-            [.., last] => last,
-        };
-
-        let magic_trailing_comma = use_magic_trailing_comma(f, last.range());
-
-        if magic_trailing_comma {
-            // A magic trailing comma forces us to print in expanded mode since we have more than
-            // one element
-            write!(
-                f,
-                [
-                    // An expanded group always needs parentheses
-                    &text("("),
-                    block_indent(&ExprSequence::new(elts)),
-                    &text(")"),
-                ]
-            )?;
-        } else if is_parenthesized(*range, elts, f)
-            && self.parentheses != TupleParentheses::StripInsideForLoop
-        {
             // If the tuple has parentheses, we generally want to keep them. The exception are for
             // loops, see `TupleParentheses::StripInsideForLoop` doc comment.
             //
             // Unlike other expression parentheses, tuple parentheses are part of the range of the
             // tuple itself.
-            write!(
+            elts if is_parenthesized(*range, elts, f)
-                f,
+                && self.parentheses != TupleParentheses::StripInsideForLoop =>
-                [group(&format_args![
+            {
-                    // If there were previously parentheses, keep them
+                write!(
-                    &text("("),
+                    f,
-                    soft_block_indent(&ExprSequence::new(elts)),
+                    [group(&format_args![
-                    &text(")"),
+                        // If there were previously parentheses, keep them
-                ])]
+                        &text("("),
-            )?;
+                        soft_block_indent(&ExprSequence::new(elts)),
-        } else {
+                        &text(")"),
-            write!(
+                    ])]
-                f,
+                )
-                [group(&format_args![
+            }
-                    // If there were previously no parentheses, add them only if the group breaks
+            elts => {
-                    if_group_breaks(&text("(")),
+                write!(
-                    soft_block_indent(&ExprSequence::new(elts)),
+                    f,
-                    if_group_breaks(&text(")")),
+                    [group(&format_args![
-                ])]
+                        // If there were previously no parentheses, add them only if the group breaks
-            )?;
+                        if_group_breaks(&text("(")),
+                        soft_block_indent(&ExprSequence::new(elts)),
+                        if_group_breaks(&text(")")),
+                    ])]
+                )
+            }
         }
-
-        Ok(())
     }
 
     fn fmt_dangling_comments(&self, _node: &ExprTuple, _f: &mut PyFormatter) -> FormatResult<()> {
@@ -152,11 +133,7 @@ impl<'a> ExprSequence<'a> {
 
 impl Format<PyFormatContext<'_>> for ExprSequence<'_> {
     fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
-        f.join_with(&format_args!(text(","), soft_line_break_or_space()))
+        f.join_comma_separated().nodes(self.elts.iter()).finish()
-            .entries(self.elts.iter().formatted())
-            .finish()?;
-        // Black style has a trailing comma on the last entry of an expanded group
-        write!(f, [if_group_breaks(&text(","))])
     }
 }
 

crates/ruff_python_formatter/src/statement/stmt_class_def.rs

@@ -1,11 +1,10 @@
-use crate::builders::use_magic_trailing_comma;
 use crate::comments::trailing_comments;
 use crate::expression::parentheses::Parenthesize;
 use crate::prelude::*;
 use crate::trivia::{SimpleTokenizer, TokenKind};
 use ruff_formatter::{format_args, write};
 use ruff_text_size::TextRange;
-use rustpython_parser::ast::{Expr, Keyword, Ranged, StmtClassDef};
+use rustpython_parser::ast::{Ranged, StmtClassDef};
 
 #[derive(Default)]
 pub struct FormatStmtClassDef;
@@ -80,10 +79,9 @@ impl Format<PyFormatContext<'_>> for FormatInheritanceClause<'_> {
             ..
         } = self.class_definition;
 
-        let separator = format_with(|f| write!(f, [text(","), soft_line_break_or_space()]));
         let source = f.context().contents();
 
-        let mut joiner = f.join_with(&separator);
+        let mut joiner = f.join_comma_separated();
 
         if let Some((first, rest)) = bases.split_first() {
             // Manually handle parentheses for the first expression because the logic in `FormatExpr`
@@ -107,23 +105,10 @@ impl Format<PyFormatContext<'_>> for FormatInheritanceClause<'_> {
                 Parenthesize::Never
             };
 
-            joiner.entry(&first.format().with_options(parenthesize));
+            joiner.entry(first, &first.format().with_options(parenthesize));
-            joiner.entries(rest.iter().formatted());
+            joiner.nodes(rest.iter());
         }
 
-        joiner.entries(keywords.iter().formatted()).finish()?;
+        joiner.nodes(keywords.iter()).finish()
-
-        if_group_breaks(&text(",")).fmt(f)?;
-
-        let last = keywords
-            .last()
-            .map(Keyword::range)
-            .or_else(|| bases.last().map(Expr::range))
-            .unwrap();
-        if use_magic_trailing_comma(f, last) {
-            hard_line_break().fmt(f)?;
-        }
-
-        Ok(())
     }
 }