Use normalize_comment

crates/ruff_python_formatter/src/comments/format.rs

@@ -1,6 +1,8 @@
-use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
+use std::borrow::Cow;
 
-use ruff_formatter::{format_args, write, FormatError, FormatState, SourceCode, VecBuffer};
+use ruff_text_size::{Ranged, TextLen, TextSize};
+
+use ruff_formatter::{format_args, write, FormatContext, FormatError, SourceCode};
 use ruff_python_ast::node::{AnyNodeRef, AstNode};
 use ruff_python_trivia::{lines_after, lines_after_ignoring_trivia, lines_before};
 
@@ -167,12 +169,17 @@ impl Format<PyFormatContext<'_>> for FormatTrailingComments<'_> {
                 //     "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
                 // )  # Some comment
                 // ```
+
+                let normalized_comment = normalize_comment(trailing, f.context().source_code())?;
+                // SAFETY: Ruff only supports formatting files <= 4GB
+                #[allow(clippy::cast_possible_truncation)]
+                let normalized_comment_len = normalized_comment.len() as u32;
                 write!(
                     f,
                     [
                         line_suffix(
                             &format_args![space(), space(), format_comment(trailing)],
-                            measure_comment(trailing, f.context())? + 2 // Account for two added spaces
+                            normalized_comment_len + 2 // Account for two added spaces
                         ),
                         expand_parent()
                     ]
@@ -275,13 +282,17 @@ impl Format<PyFormatContext<'_>> for FormatDanglingOpenParenthesisComments<'_> {
                 "Expected dangling comment to be at the end of the line"
             );
 
+            let normalized_comment = normalize_comment(comment, f.context().source_code())?;
+            // SAFETY: Ruff only supports formatting files <= 4GB
+            #[allow(clippy::cast_possible_truncation)]
+            let normalized_comment_len = normalized_comment.len() as u32;
             write!(
                 f,
                 [
                     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.
-                        measure_comment(comment, f.context())? + 2 // Account for two added spaces
+                        normalized_comment_len + 2 // Account for two added spaces
                     ),
                     expand_parent()
                 ]
@@ -307,9 +318,10 @@ pub(crate) struct FormatComment<'a> {
 
 impl Format<PyFormatContext<'_>> for FormatComment<'_> {
     fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
-        // We don't need the formatted comment's width.
-        let _ = write_comment(f, self.comment)?;
-        Ok(())
+        let source = SourceCode::new(f.context().source());
+        let comment = normalize_comment(self.comment, source)?;
+
+        dynamic_text(&comment, None).fmt(f)
     }
 }
 
@@ -348,24 +360,13 @@ impl Format<PyFormatContext<'_>> for FormatEmptyLines {
     }
 }
 
-/// A helper used to measure formatted comments.
-///
-/// Use a temporary formatter to write a normalized, formatted comment
-/// to in order to compute its width for a reserved-width line suffix element.
-fn measure_comment(comment: &SourceComment, context: &PyFormatContext) -> FormatResult<u32> {
-    let mut state = FormatState::new(context.clone());
-    let mut buffer = VecBuffer::new(&mut state);
-    let comment_len = write_comment(&mut Formatter::new(&mut buffer), comment)?;
-    Ok(comment_len)
-}
-
-/// Write a comment to a formatter and return the normalized comment's width.
-fn write_comment(f: &mut PyFormatter, comment: &SourceComment) -> FormatResult<u32> {
+/// A helper for normalizing comments efficiently.
+fn normalize_comment<'a>(
+    comment: &'a SourceComment,
+    source: SourceCode<'a>,
+) -> FormatResult<Cow<'a, str>> {
     let slice = comment.slice();
-    let comment_text = slice.text(SourceCode::new(f.context().source()));
-
-    // Track any additional width the formatted comment will have after normalization.
-    let mut added_width = TextSize::new(0);
+    let comment_text = slice.text(source);
 
     let trimmed = comment_text.trim_end();
     let trailing_whitespace_len = comment_text.text_len() - trimmed.text_len();
@@ -376,41 +377,26 @@ fn write_comment(f: &mut PyFormatter, comment: &SourceComment) -> FormatResult<u
         ));
     };
 
-    // Fast path for correctly formatted comments:
-    // * Start with a `#` and are followed by a space
-    // * Have no trailing whitespace.
-    if trailing_whitespace_len == TextSize::new(0) && content.starts_with(' ') {
-        source_text_slice(slice.range(), ContainsNewlines::No).fmt(f)?;
-        return Ok(slice.range().len().into());
+    if trailing_whitespace_len == TextSize::new(0) {
+        // Fast path for correctly formatted comments:
+        // * Start with a `# '.
+        // * Have no trailing whitespace.
+        if content.starts_with(' ') {
+            return Ok(Cow::Borrowed(comment_text));
+        }
     }
 
-    write!(f, [source_position(slice.start()), text("#")])?;
-
-    // Starts with a non breaking space
-    let start_offset =
-        if content.starts_with('\u{A0}') && !content.trim_start().starts_with("type:") {
-            // Replace non-breaking space with a space (if not followed by a normal space)
-            "#\u{A0}".text_len()
-        } else {
-            '#'.text_len()
-        };
-
-    // Add a space between the `#` and the text if the source contains none.
+    // Formatted comments start with '# '. We perform this normalization if it's necessary.
     if !content.is_empty() && !content.starts_with([' ', '!', ':', '#', '\'']) {
-        write!(f, [space()])?;
-        added_width += TextSize::new(1);
+        // Only replace non-breaking space with a space if it's not followed by a normal space.
+        if !(content.starts_with('\u{A0}')
+            && !content.trim_start().starts_with("type:")
+            && !content.trim_start_matches('\u{A0}').starts_with(' '))
+        {
+            return Ok(Cow::Owned(std::format!("# {}", content.trim_start())));
+        }
     }
 
-    let start = slice.start() + start_offset;
-    let end = slice.end() - trailing_whitespace_len;
-
-    write!(
-        f,
-        [
-            source_text_slice(TextRange::new(start, end), ContainsNewlines::No),
-            source_position(slice.end())
-        ]
-    )?;
-
-    Ok((end - slice.start() + added_width).into())
+    // Finally we return an new, owned comment with a leading '#' since we removed it earlier.
+    Ok(Cow::Owned(std::format!("#{content}")))
 }