diff --git a/crates/ruff_python_formatter/src/comments/format.rs b/crates/ruff_python_formatter/src/comments/format.rs index 26efa692864b80..79e0aff9e7b73f 100644 --- a/crates/ruff_python_formatter/src/comments/format.rs +++ b/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> 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> 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> 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> 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 { - 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 { +/// A helper for normalizing comments efficiently. +fn normalize_comment<'a>( + comment: &'a SourceComment, + source: SourceCode<'a>, +) -> FormatResult> { 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