Use reserved width to include line suffix measurement

crates/ruff_formatter/src/builders.rs

@@ -435,46 +435,78 @@ fn debug_assert_no_newlines(text: &str) {
     debug_assert!(!text.contains('\r'), "The content '{text}' contains an unsupported '\\r' line terminator character but text must only use line feeds '\\n' as line separator. Use '\\n' instead of '\\r' and '\\r\\n' to insert a line break in strings.");
 }
 
-/// Pushes some content to the end of the current line
+/// Pushes some content to the end of the current line.
 ///
 /// ## Examples
 ///
-/// ```
-/// use ruff_formatter::{format};
+/// ```rust
+/// use ruff_formatter::format;
 /// use ruff_formatter::prelude::*;
 ///
-/// fn main() -> FormatResult<()> {
+/// # fn main() -> FormatResult<()> {
 /// let elements = format!(SimpleFormatContext::default(), [
 ///     text("a"),
-///     line_suffix(&text("c")),
+///     line_suffix(&text("c"), 0),
 ///     text("b")
 /// ])?;
 ///
-/// assert_eq!(
-///     "abc",
-///     elements.print()?.as_code()
-/// );
+/// assert_eq!("abc", elements.print()?.as_code());
+/// # Ok(())
+/// # }
+/// ```
+///
+/// Provide reserved width for the line suffix to include it during measurement.
+/// ```rust
+/// use ruff_formatter::{format, format_args, LineWidth, SimpleFormatContext, SimpleFormatOptions};
+/// use ruff_formatter::prelude::*;
+///
+/// # fn main() -> FormatResult<()> {
+/// let context = SimpleFormatContext::new(SimpleFormatOptions {
+///     line_width: LineWidth::try_from(10).unwrap(),
+///     ..SimpleFormatOptions::default()
+/// });
+///
+/// let elements = format!(context, [
+///     // Breaks
+///     group(&format_args![
+///         if_group_breaks(&text("(")),
+///         soft_block_indent(&format_args![text("a"), line_suffix(&text(" // a comment"), 13)]),
+///         if_group_breaks(&text(")"))
+///         ]),
+///
+///     // Fits
+///     group(&format_args![
+///         if_group_breaks(&text("(")),
+///         soft_block_indent(&format_args![text("a"), line_suffix(&text(" // a comment"), 0)]),
+///         if_group_breaks(&text(")"))
+///     ]),
+/// ])?;
+/// # assert_eq!("(\n\ta // a comment\n)a // a comment", elements.print()?.as_code());
 /// # Ok(())
 /// # }
 /// ```
 #[inline]
-pub fn line_suffix<Content, Context>(inner: &Content) -> LineSuffix<Context>
+pub fn line_suffix<Content, Context>(inner: &Content, reserved_width: u32) -> LineSuffix<Context>
 where
     Content: Format<Context>,
 {
     LineSuffix {
         content: Argument::new(inner),
+        reserved_width,
     }
 }
 
 #[derive(Copy, Clone)]
 pub struct LineSuffix<'a, Context> {
     content: Argument<'a, Context>,
+    reserved_width: u32,
 }
 
 impl<Context> Format<Context> for LineSuffix<'_, Context> {
     fn fmt(&self, f: &mut Formatter<Context>) -> FormatResult<()> {
-        f.write_element(FormatElement::Tag(StartLineSuffix));
+        f.write_element(FormatElement::Tag(StartLineSuffix {
+            reserved_width: self.reserved_width,
+        }));
         Arguments::from(&self.content).fmt(f)?;
         f.write_element(FormatElement::Tag(EndLineSuffix));
 
@@ -501,7 +533,7 @@ impl<Context> std::fmt::Debug for LineSuffix<'_, Context> {
 /// # fn  main() -> FormatResult<()> {
 /// let elements = format!(SimpleFormatContext::default(), [
 ///     text("a"),
-///     line_suffix(&text("c")),
+///     line_suffix(&text("c"), 0),
 ///     text("b"),
 ///     line_suffix_boundary(),
 ///     text("d")

crates/ruff_formatter/src/format_element/document.rs

@@ -459,8 +459,16 @@ impl Format<IrFormatContext<'_>> for &[FormatElement] {
                             )?;
                         }
 
-                        StartLineSuffix => {
-                            write!(f, [text("line_suffix(")])?;
+                        StartLineSuffix { reserved_width } => {
+                            write!(
+                                f,
+                                [
+                                    text("line_suffix("),
+                                    dynamic_text(&std::format!("{reserved_width:?}"), None),
+                                    text(","),
+                                    space(),
+                                ]
+                            )?;
                         }
 
                         StartVerbatim(_) => {
@@ -672,7 +680,9 @@ impl FormatElements for [FormatElement] {
             match element {
                 // Line suffix
                 // Ignore if any of its content breaks
-                FormatElement::Tag(Tag::StartLineSuffix | Tag::StartFitsExpanded(_)) => {
+                FormatElement::Tag(
+                    Tag::StartLineSuffix { reserved_width: _ } | Tag::StartFitsExpanded(_),
+                ) => {
                     ignore_depth += 1;
                 }
                 FormatElement::Tag(Tag::EndLineSuffix | Tag::EndFitsExpanded) => {

crates/ruff_formatter/src/format_element/tag.rs

@@ -63,8 +63,11 @@ pub enum Tag {
     StartEntry,
     EndEntry,
 
-    /// Delay the printing of its content until the next line break
-    StartLineSuffix,
+    /// Delay the printing of its content until the next line break. Using reserved width will include
+    /// the associated line suffix during measurement.
+    StartLineSuffix {
+        reserved_width: u32,
+    },
     EndLineSuffix,
 
     /// A token that tracks tokens/nodes that are printed as verbatim.
@@ -96,7 +99,7 @@ impl Tag {
                 | Tag::StartIndentIfGroupBreaks(_)
                 | Tag::StartFill
                 | Tag::StartEntry
-                | Tag::StartLineSuffix
+                | Tag::StartLineSuffix { reserved_width: _ }
                 | Tag::StartVerbatim(_)
                 | Tag::StartLabelled(_)
                 | Tag::StartFitsExpanded(_)
@@ -122,7 +125,7 @@ impl Tag {
             StartIndentIfGroupBreaks(_) | EndIndentIfGroupBreaks => TagKind::IndentIfGroupBreaks,
             StartFill | EndFill => TagKind::Fill,
             StartEntry | EndEntry => TagKind::Entry,
-            StartLineSuffix | EndLineSuffix => TagKind::LineSuffix,
+            StartLineSuffix { reserved_width: _ } | EndLineSuffix => TagKind::LineSuffix,
             StartVerbatim(_) | EndVerbatim => TagKind::Verbatim,
             StartLabelled(_) | EndLabelled => TagKind::Labelled,
             StartFitsExpanded { .. } | EndFitsExpanded => TagKind::FitsExpanded,

crates/ruff_formatter/src/printer/mod.rs

@@ -233,7 +233,8 @@ impl<'a> Printer<'a> {
                 stack.push(TagKind::IndentIfGroupBreaks, args);
             }
 
-            FormatElement::Tag(StartLineSuffix) => {
+            FormatElement::Tag(StartLineSuffix { reserved_width }) => {
+                self.state.line_width += reserved_width;
                 self.state
                     .line_suffixes
                     .extend(args, queue.iter_content(TagKind::LineSuffix));
@@ -1188,7 +1189,11 @@ impl<'a, 'print> FitsMeasurer<'a, 'print> {
                 }
             }
 
-            FormatElement::Tag(StartLineSuffix) => {
+            FormatElement::Tag(StartLineSuffix { reserved_width }) => {
+                self.state.line_width += reserved_width;
+                if self.state.line_width > self.options().print_width.into() {
+                    return Ok(Fits::No);
+                }
                 self.queue.skip_content(TagKind::LineSuffix);
                 self.state.has_line_suffix = true;
             }
@@ -1724,12 +1729,42 @@ two lines`,
                 text("]")
             ]),
             text(";"),
-            line_suffix(&format_args![space(), text("// trailing")])
+            line_suffix(&format_args![space(), text("// trailing")], 0)
         ]);
 
         assert_eq!(printed.as_code(), "[1, 2, 3]; // trailing");
     }
 
+    #[test]
+    fn line_suffix_with_reserved_width() {
+        let printed = format(&format_args![
+            group(&format_args![
+                text("["),
+                soft_block_indent(&format_with(|f| {
+                    f.fill()
+                        .entry(
+                            &soft_line_break_or_space(),
+                            &format_args!(text("1"), text(",")),
+                        )
+                        .entry(
+                            &soft_line_break_or_space(),
+                            &format_args!(text("2"), text(",")),
+                        )
+                        .entry(
+                            &soft_line_break_or_space(),
+                            &format_args!(text("3"), if_group_breaks(&text(","))),
+                        )
+                        .finish()
+                })),
+                text("]")
+            ]),
+            text(";"),
+            line_suffix(&format_args![space(), text("// Using reserved width causes this content to not fit even though it's a line suffix element")], 93)
+        ]);
+
+        assert_eq!(printed.as_code(), "[\n  1, 2, 3\n]; // Using reserved width causes this content to not fit even though it's a line suffix element");
+    }
+
     #[test]
     fn conditional_with_group_id_in_fits() {
         let content = format_with(|f| {

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

@@ -66,3 +66,5 @@
 h1 = ((((1, 2))))
 h2 = ((((1, "qweiurpoiqwurepqiurpqirpuqoiwrupqoirupqoirupqoiurpqiorupwqiourpqurpqurpqurpqurpqurpqurüqurqpuriq"))))
 h3 = 1, "qweiurpoiqwurepqiurpqirpuqoiwrupqoirupqoirupqoiurpqiorupwqiourpqurpqurpqurpqurpqurpqurüqurqpuriq"
+
+i1 = ("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",)  # This should break

crates/ruff_python_formatter/src/comments/format.rs

@@ -1,3 +1,5 @@
+use std::borrow::Cow;
+
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 
 use ruff_formatter::{format_args, write, FormatError, SourceCode};
@@ -150,18 +152,34 @@ impl Format<PyFormatContext<'_>> for FormatTrailingComments<'_> {
                 write!(
                     f,
                     [
-                        line_suffix(&format_args![
-                            empty_lines(lines_before_comment),
-                            format_comment(trailing)
-                        ]),
+                        line_suffix(
+                            &format_args![
+                                empty_lines(lines_before_comment),
+                                format_comment(trailing)
+                            ],
+                            0 // Reserving width isn't necessary because we don't split comments and the empty lines expand any enclosing group.
+                        ),
                         expand_parent()
                     ]
                 )?;
             } else {
+                // A trailing comment at the end of a line has a reserved width to consider during line measurement.
+                // ```python
+                // tup = (
+                //     "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+                // )  # Some comment
+                // ```
+                let source = SourceCode::new(f.context().source());
+                // SAFETY: Ruff only supports formatting files <= 4GB
+                #[allow(clippy::cast_possible_truncation)]
+                let comment_len = normalize_comment(trailing, source).len() as u32;
                 write!(
                     f,
                     [
-                        line_suffix(&format_args![space(), space(), format_comment(trailing)]),
+                        line_suffix(
+                            &format_args![space(), space(), format_comment(trailing)],
+                            comment_len + 2 // Account for two added spaces
+                        ),
                         expand_parent()
                     ]
                 )?;
@@ -263,10 +281,18 @@ impl Format<PyFormatContext<'_>> for FormatDanglingOpenParenthesisComments<'_> {
                 "Expected dangling comment to be at the end of the line"
             );
 
+            let source = SourceCode::new(f.context().source());
+            // SAFETY: Ruff only supports formatting files <= 4GB
+            #[allow(clippy::cast_possible_truncation)]
+            let comment_len = normalize_comment(comment, source).len() as u32;
             write!(
                 f,
                 [
-                    line_suffix(&format_args!(space(), space(), format_comment(comment))),
+                    line_suffix(
+                        &format_args!(space(), space(), format_comment(comment)),
+                        // Marking the comment as a line suffix with reserved width is safe since we expect the comment to be end of line.
+                        comment_len + 2 // Account for two added spaces
+                    ),
                     expand_parent()
                 ]
             )?;
@@ -373,3 +399,44 @@ impl Format<PyFormatContext<'_>> for FormatEmptyLines {
         }
     }
 }
+
+// TODO(cnpryer): Alternatively could just duplicate the range tracking from normalization instead
+//   of costly String operations and allocations.
+//   Could abstract the comment behind `NormalizedComment`
+fn normalize_comment<'a>(comment: &'a SourceComment, source: SourceCode<'a>) -> Cow<'a, str> {
+    let slice = comment.slice();
+    let comment_text = slice.text(source);
+
+    // If the comment isn't leading with '#' it's invalid, so we return the text as-is borrowed.
+    // TODO(cnpryer): If this is always used alongside `.fmt()` on comments this should be safe, but
+    //   we can probably do better here.
+    let Some(content) = comment_text.strip_prefix('#') else {
+        return Cow::Borrowed(comment_text);
+    };
+
+    // Make the content mutable in case we need to normalize it. Any normalization is done with lazy operations.
+    let mut content = Cow::Borrowed(content);
+
+    // If the comment has trailing whitespace then we take ownership of a string to make mutations to
+    let trimmed = content.trim_end();
+    let trailing_whitespace_len = content.text_len() - trimmed.text_len();
+    if trailing_whitespace_len > TextSize::new(0) {
+        content = Cow::Owned(trimmed.to_owned());
+    }
+
+    // Fast path for correctly formatted comments:
+    // * Start with a `# '.
+    // * Have no trailing whitespace.
+    if trailing_whitespace_len == TextSize::new(0) && content.starts_with(' ') {
+        return Cow::Borrowed(comment_text);
+    }
+
+    // Formatted comments start with '# '. We perform this normalization if it's necessary.
+    if !content.is_empty() && !content.starts_with([' ', '!', ':', '#', '\'', '\u{A0}']) {
+        Cow::Owned(std::format!("# {content}"))
+    } else if trailing_whitespace_len > TextSize::new(0) {
+        Cow::Owned(std::format!("#{content}"))
+    } else {
+        Cow::Borrowed(comment_text)
+    }
+}

crates/ruff_python_formatter/src/expression/mod.rs

@@ -247,10 +247,13 @@ impl Format<PyFormatContext<'_>> for MaybeParenthesizeExpression<'_> {
                     if format_expression.inspect(f)?.will_break() {
                         // The group here is necessary because `format_expression` may contain IR elements
                         // that refer to the group id
-                        group(&format_expression)
-                            .with_group_id(Some(group_id))
-                            .should_expand(true)
-                            .fmt(f)
+                        group(&format_args![
+                            text("("),
+                            soft_block_indent(&format_expression),
+                            text(")")
+                        ])
+                        .with_group_id(Some(group_id))
+                        .fmt(f)
                     } else {
                         // Only add parentheses if it makes the expression fit on the line.
                         // Using the flat version as the most expanded version gives a left-to-right splitting behavior

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

@@ -156,7 +156,7 @@ aaaaaaaaaaaaa, bbbbbbbbb = map(list, map(itertools.chain.from_iterable, zip(*ite
      )
 
 
-@@ -108,11 +112,18 @@
+@@ -108,11 +112,20 @@
      "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
  )
 
@@ -176,7 +176,10 @@ aaaaaaaaaaaaa, bbbbbbbbb = map(list, map(itertools.chain.from_iterable, zip(*ite
 +    ],  # type: ignore
  )
 
- aaaaaaaaaaaaa, bbbbbbbbb = map(list, map(itertools.chain.from_iterable, zip(*items)))  # type: ignore[arg-type]
+-aaaaaaaaaaaaa, bbbbbbbbb = map(list, map(itertools.chain.from_iterable, zip(*items)))  # type: ignore[arg-type]
++aaaaaaaaaaaaa, bbbbbbbbb = map(
++    list, map(itertools.chain.from_iterable, zip(*items))
++)  # type: ignore[arg-type]
 ```
 
 ## Ruff Output
@@ -310,7 +313,9 @@ call_to_some_function_asdf(
     ],  # type: ignore
 )
 
-aaaaaaaaaaaaa, bbbbbbbbb = map(list, map(itertools.chain.from_iterable, zip(*items)))  # type: ignore[arg-type]
+aaaaaaaaaaaaa, bbbbbbbbb = map(
+    list, map(itertools.chain.from_iterable, zip(*items))
+)  # type: ignore[arg-type]
 ```
 
 ## Black Output