Add reserved_width to LineSuffix

astral-sh · Aug 26, 2023 · 5de6854 · 5de6854
1 parent 15b73bd
commit 5de6854
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 20 deletions.
diff --git a/crates/ruff_formatter/src/builders.rs b/crates/ruff_formatter/src/builders.rs
@@ -434,7 +434,8 @@ 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. Provide a reserved width in
+/// order to include the line suffix content during measurement.
 ///
 /// ## Examples
 ///
@@ -445,7 +446,7 @@ fn debug_assert_no_newlines(text: &str) {
 /// fn main() -> FormatResult<()> {
 /// let elements = format!(SimpleFormatContext::default(), [
 ///     text("a"),
-///     line_suffix(&text("c")),
+///     line_suffix(&text("c"), 0),
 ///     text("b")
 /// ])?;
 ///
@@ -457,23 +458,25 @@ fn debug_assert_no_newlines(text: &str) {
 /// # }
 /// ```
 #[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(self.reserved_width)));
         Arguments::from(&self.content).fmt(f)?;
         f.write_element(FormatElement::Tag(EndLineSuffix));
 
@@ -500,7 +503,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")

diff --git a/crates/ruff_formatter/src/format_element/document.rs b/crates/ruff_formatter/src/format_element/document.rs
@@ -458,8 +458,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(_) => {
@@ -671,7 +679,7 @@ 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(_) | Tag::StartFitsExpanded(_)) => {
                     ignore_depth += 1;
                 }
                 FormatElement::Tag(Tag::EndLineSuffix | Tag::EndFitsExpanded) => {

diff --git a/crates/ruff_formatter/src/format_element/tag.rs b/crates/ruff_formatter/src/format_element/tag.rs
@@ -64,7 +64,7 @@ pub enum Tag {
     EndEntry,
 
     /// Delay the printing of its content until the next line break
-    StartLineSuffix,
+    StartLineSuffix(u32),
     EndLineSuffix,
 
     /// A token that tracks tokens/nodes that are printed as verbatim.
@@ -96,7 +96,7 @@ impl Tag {
                 | Tag::StartIndentIfGroupBreaks(_)
                 | Tag::StartFill
                 | Tag::StartEntry
-                | Tag::StartLineSuffix
+                | Tag::StartLineSuffix(_)
                 | Tag::StartVerbatim(_)
                 | Tag::StartLabelled(_)
                 | Tag::StartFitsExpanded(_)
@@ -122,7 +122,7 @@ impl Tag {
             StartIndentIfGroupBreaks(_) | EndIndentIfGroupBreaks => TagKind::IndentIfGroupBreaks,
             StartFill | EndFill => TagKind::Fill,
             StartEntry | EndEntry => TagKind::Entry,
-            StartLineSuffix | EndLineSuffix => TagKind::LineSuffix,
+            StartLineSuffix(_) | EndLineSuffix => TagKind::LineSuffix,
             StartVerbatim(_) | EndVerbatim => TagKind::Verbatim,
             StartLabelled(_) | EndLabelled => TagKind::Labelled,
             StartFitsExpanded { .. } | EndFitsExpanded => TagKind::FitsExpanded,

diff --git a/crates/ruff_formatter/src/printer/mod.rs b/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));
@@ -1184,7 +1185,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;
             }
@@ -1720,7 +1725,7 @@ 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");

diff --git a/crates/ruff_python_formatter/src/comments/format.rs b/crates/ruff_python_formatter/src/comments/format.rs
@@ -151,18 +151,21 @@ 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
+                        ),
                         expand_parent()
                     ]
                 )?;
             } else {
                 write!(
                     f,
                     [
-                        line_suffix(&format_args![space(), space(), format_comment(trailing)]),
+                        line_suffix(&format_args![space(), space(), format_comment(trailing)], 0),
                         expand_parent()
                     ]
                 )?;
@@ -267,7 +270,7 @@ impl Format<PyFormatContext<'_>> for FormatDanglingOpenParenthesisComments<'_> {
             write!(
                 f,
                 [
-                    line_suffix(&format_args!(space(), space(), format_comment(comment))),
+                    line_suffix(&format_args!(space(), space(), format_comment(comment)), 0),
                     expand_parent()
                 ]
             )?;