diff --git a/crates/ruff_python_formatter/src/comments/mod.rs b/crates/ruff_python_formatter/src/comments/mod.rs index 91ff7e2b8b..ab408fc124 100644 --- a/crates/ruff_python_formatter/src/comments/mod.rs +++ b/crates/ruff_python_formatter/src/comments/mod.rs @@ -481,6 +481,28 @@ impl<'a> Comments<'a> { pub(crate) type LeadingDanglingTrailingComments<'a> = LeadingDanglingTrailing<'a, SourceComment>; +impl LeadingDanglingTrailingComments<'_> { + /// Returns `true` if the struct has any [leading comments](self#leading-comments). + #[inline] + pub(crate) fn has_leading(&self) -> bool { + !self.leading.is_empty() + } + + /// Returns `true` if the struct has any [trailing comments](self#trailing-comments). + #[inline] + pub(crate) fn has_trailing(&self) -> bool { + !self.trailing.is_empty() + } + + /// Returns `true` if the struct has any [trailing own line comments](self#trailing-comments). + #[inline] + pub(crate) fn has_trailing_own_line(&self) -> bool { + self.trailing + .iter() + .any(|comment| comment.line_position().is_own_line()) + } +} + #[derive(Debug, Default)] struct CommentsData<'a> { comments: CommentsMap<'a>, diff --git a/crates/ruff_python_formatter/src/statement/suite.rs b/crates/ruff_python_formatter/src/statement/suite.rs index 6adbe658ca..de5ab8e4da 100644 --- a/crates/ruff_python_formatter/src/statement/suite.rs +++ b/crates/ruff_python_formatter/src/statement/suite.rs @@ -142,18 +142,22 @@ impl FormatRule> for FormatSuite { ) }; + let mut preceding_comments = comments.leading_dangling_trailing(preceding); + while let Some(following) = iter.next() { + let following_comments = comments.leading_dangling_trailing(following); + // Add empty lines before and after a function or class definition. If the preceding // node is a function or class, and contains trailing comments, then the statement // itself will add the requisite empty lines when formatting its comments. if (is_class_or_function_definition(preceding) - && !comments.has_trailing_own_line(preceding)) + && !preceding_comments.has_trailing_own_line()) || is_class_or_function_definition(following) { match self.kind { SuiteKind::TopLevel if source_type.is_stub() => { // Preserve the empty line if the definitions are separated by a comment - if comments.has_trailing(preceding) || comments.has_leading(following) { + if preceding_comments.has_trailing() || following_comments.has_leading() { empty_line().fmt(f)?; } else { // Two subsequent classes that both have an ellipsis only body @@ -196,7 +200,7 @@ impl FormatRule> for FormatSuite { } } } else if is_import_definition(preceding) - && (!is_import_definition(following) || comments.has_leading(following)) + && (!is_import_definition(following) || following_comments.has_leading()) { // Enforce _at least_ one empty line after an import statement (but allow up to // two at the top-level). In this context, "after an import statement" means that @@ -230,7 +234,7 @@ impl FormatRule> for FormatSuite { // which is 0 instead of 1, the number of lines between the trailing comment and // the leading comment. This is why the suite handling counts the lines before the // start of the next statement or before the first leading comments for compound statements. - let start = if let Some(first_leading) = comments.leading(following).first() { + let start = if let Some(first_leading) = following_comments.leading.first() { first_leading.start() } else { following.start() @@ -285,8 +289,8 @@ impl FormatRule> for FormatSuite { lines_after(offset, source) }; - let end = comments - .trailing(preceding) + let end = preceding_comments + .trailing .last() .map_or(preceding.end(), |comment| comment.slice().end()); @@ -306,8 +310,6 @@ impl FormatRule> for FormatSuite { } } - let following_comments = comments.leading_dangling_trailing(following); - if following_comments .leading .iter() @@ -318,6 +320,7 @@ impl FormatRule> for FormatSuite { &mut iter, f, )?; + preceding_comments = comments.leading_dangling_trailing(preceding); } else if following_comments .trailing .iter() @@ -328,10 +331,13 @@ impl FormatRule> for FormatSuite { &mut iter, f, )?; + preceding_comments = comments.leading_dangling_trailing(preceding); } else { following.format().fmt(f)?; preceding = following; + preceding_comments = following_comments; } + after_class_docstring = false; }