Format while Statement (#4810)

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/while.py

@@ -0,0 +1,30 @@
+while 34: # trailing test comment
+ pass # trailing last statement comment
+
+ # trailing while body comment
+
+# leading else comment
+
+else: # trailing else comment
+ pass
+
+ # trailing else body comment
+
+
+while aVeryLongConditionThatSpillsOverToTheNextLineBecauseItIsExtremelyLongAndGoesOnAndOnAndOnAndOnAndOnAndOnAndOnAndOnAndOn: # trailing comment
+ pass
+
+else:
+ ...
+
+while (
+ some_condition(unformatted, args) and anotherCondition or aThirdCondition
+): # comment
+ print("Do something")
+
+
+while (
+ some_condition(unformatted, args) # trailing some condition
+ and anotherCondition or aThirdCondition # trailing third condition
+): # comment
+ print("Do something")

crates/ruff_python_formatter/src/builders.rs

@@ -11,7 +11,7 @@ pub(crate) trait PyFormatterExtensions<'ast, 'buf> {
  /// empty lines between any two nodes. Separates any two nodes by at least a hard line break.
  ///
  /// * [`NodeLevel::Module`]: Up to two empty lines
- /// * [`NodeLevel::Statement`]: Up to one empty line
+ /// * [`NodeLevel::CompoundStatement`]: Up to one empty line
  /// * [`NodeLevel::Parenthesized`]: No empty lines
  fn join_nodes<'fmt>(&'fmt mut self, level: NodeLevel) -> JoinNodesBuilder<'fmt, 'ast, 'buf>;
 }
@@ -53,10 +53,12 @@ impl<'fmt, 'ast, 'buf> JoinNodesBuilder<'fmt, 'ast, 'buf> {
  2 => empty_line().fmt(f),
  _ => write!(f, [empty_line(), empty_line()]),
  },
- NodeLevel::Statement => match lines_before(f.context().contents(), node.start()) {
- 0 | 1 => hard_line_break().fmt(f),
- _ => empty_line().fmt(f),
- },
+ NodeLevel::CompoundStatement => {
+ match lines_before(f.context().contents(), node.start()) {
+ 0 | 1 => hard_line_break().fmt(f),
+ _ => empty_line().fmt(f),
+ }
+ }
  NodeLevel::Parenthesized => hard_line_break().fmt(f),
  });
 
@@ -180,7 +182,7 @@ no_leading_newline = 30"#
  // Should keep at most one empty level
  #[test]
  fn ranged_builder_statement_level() {
- let printed = format_ranged(NodeLevel::Statement);
+ let printed = format_ranged(NodeLevel::CompoundStatement);
 
  assert_eq!(
  &printed,

crates/ruff_python_formatter/src/comments/format.rs

@@ -6,27 +6,37 @@ use ruff_formatter::{format_args, write, FormatError, SourceCode};
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::prelude::AstNode;
 use ruff_text_size::{TextLen, TextRange, TextSize};
+use rustpython_parser::ast::Ranged;
 
 /// Formats the leading comments of a node.
-pub(crate) fn leading_comments<T>(node: &T) -> FormatLeadingComments
+pub(crate) fn leading_node_comments<T>(node: &T) -> FormatLeadingComments
 where
  T: AstNode,
 {
- FormatLeadingComments {
- node: node.as_any_node_ref(),
- }
+ FormatLeadingComments::Node(node.as_any_node_ref())
+}
+
+/// Formats the passed comments as leading comments
+pub(crate) const fn leading_comments(comments: &[SourceComment]) -> FormatLeadingComments {
+ FormatLeadingComments::Comments(comments)
 }
 
 #[derive(Copy, Clone, Debug)]
-pub(crate) struct FormatLeadingComments<'a> {
- node: AnyNodeRef<'a>,
+pub(crate) enum FormatLeadingComments<'a> {
+ Node(AnyNodeRef<'a>),
+ Comments(&'a [SourceComment]),
 }
 
 impl Format<PyFormatContext<'_>> for FormatLeadingComments<'_> {
  fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
  let comments = f.context().comments().clone();
 
- for comment in comments.leading_comments(self.node) {
+ let leading_comments = match self {
+ FormatLeadingComments::Node(node) => comments.leading_comments(*node),
+ FormatLeadingComments::Comments(comments) => comments,
+ };
+
+ for comment in leading_comments {
  let slice = comment.slice();
 
  let lines_after_comment = lines_after(f.context().contents(), slice.end());
@@ -42,32 +52,88 @@ impl Format<PyFormatContext<'_>> for FormatLeadingComments<'_> {
  }
 }
 
+/// Formats the leading `comments` of an alternate branch and ensures that it preserves the right
+/// number of empty lines before. The `last_node` is the last node of the preceding body.
+///
+/// For example, `last_node` is the last statement in the if body when formatting the leading
+/// comments of the `else` branch.
+pub(crate) fn leading_alternate_branch_comments<'a, T>(
+ comments: &'a [SourceComment],
+ last_node: Option<T>,
+) -> FormatLeadingAlternateBranchComments<'a>
+where
+ T: Into<AnyNodeRef<'a>>,
+{
+ FormatLeadingAlternateBranchComments {
+ comments,
+ last_node: last_node.map(std::convert::Into::into),
+ }
+}
+
+pub(crate) struct FormatLeadingAlternateBranchComments<'a> {
+ comments: &'a [SourceComment],
+ last_node: Option<AnyNodeRef<'a>>,
+}
+
+impl Format<PyFormatContext<'_>> for FormatLeadingAlternateBranchComments<'_> {
+ fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+ if let Some(first_leading) = self.comments.first() {
+ // Leading comments only preserves the lines after the comment but not before.
+ // Insert the necessary lines.
+ if lines_before(f.context().contents(), first_leading.slice().start()) > 1 {
+ write!(f, [empty_line()])?;
+ }
+
+ write!(f, [leading_comments(self.comments)])?;
+ } else if let Some(last_preceding) = self.last_node {
+ // The leading comments formatting ensures that it preserves the right amount of lines after
+ // We need to take care of this ourselves, if there's no leading `else` comment.
+ if lines_after(f.context().contents(), last_preceding.end()) > 1 {
+ write!(f, [empty_line()])?;
+ }
+ }
+
+ Ok(())
+ }
+}
+
 /// Formats the trailing comments of `node`
-pub(crate) fn trailing_comments<T>(node: &T) -> FormatTrailingComments
+pub(crate) fn trailing_node_comments<T>(node: &T) -> FormatTrailingComments
 where
  T: AstNode,
 {
- FormatTrailingComments {
- node: node.as_any_node_ref(),
- }
+ FormatTrailingComments::Node(node.as_any_node_ref())
 }
 
-pub(crate) struct FormatTrailingComments<'a> {
- node: AnyNodeRef<'a>,
+/// Formats the passed comments as trailing comments
+pub(crate) fn trailing_comments(comments: &[SourceComment]) -> FormatTrailingComments {
+ FormatTrailingComments::Comments(comments)
+}
+
+pub(crate) enum FormatTrailingComments<'a> {
+ Node(AnyNodeRef<'a>),
+ Comments(&'a [SourceComment]),
 }
 
 impl Format<PyFormatContext<'_>> for FormatTrailingComments<'_> {
  fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
  let comments = f.context().comments().clone();
- let mut has_empty_lines_before = false;
 
- for trailing in comments.trailing_comments(self.node) {
+ let trailing_comments = match self {
+ FormatTrailingComments::Node(node) => comments.trailing_comments(*node),
+ FormatTrailingComments::Comments(comments) => comments,
+ };
+
+ let mut has_trailing_own_line_comment = false;
+
+ for trailing in trailing_comments {
  let slice = trailing.slice();
 
- let lines_before_comment = lines_before(f.context().contents(), slice.start());
- has_empty_lines_before |= lines_before_comment > 0;
+ has_trailing_own_line_comment |= trailing.position().is_own_line();
+
+ if has_trailing_own_line_comment {
+ let lines_before_comment = lines_before(f.context().contents(), slice.start());
 
- if has_empty_lines_before {
  // A trailing comment at the end of a body or list
  // ```python
  // def test():
@@ -105,7 +171,7 @@ impl Format<PyFormatContext<'_>> for FormatTrailingComments<'_> {
 }
 
 /// Formats the dangling comments of `node`.
-pub(crate) fn dangling_comments<T>(node: &T) -> FormatDanglingComments
+pub(crate) fn dangling_node_comments<T>(node: &T) -> FormatDanglingComments
 where
  T: AstNode,
 {
@@ -229,7 +295,7 @@ impl Format<PyFormatContext<'_>> for FormatEmptyLines {
  _ => write!(f, [empty_line(), empty_line()]),
  },
 
- NodeLevel::Statement => match self.lines {
+ NodeLevel::CompoundStatement => match self.lines {
  0 | 1 => write!(f, [hard_line_break()]),
  _ => write!(f, [empty_line()]),
  },

crates/ruff_python_formatter/src/comments/mod.rs

@@ -103,7 +103,10 @@ use crate::comments::debug::{DebugComment, DebugComments};
 use crate::comments::map::MultiMap;
 use crate::comments::node_key::NodeRefEqualityKey;
 use crate::comments::visitor::CommentsVisitor;
-pub(crate) use format::{dangling_comments, leading_comments, trailing_comments};
+pub(crate) use format::{
+ dangling_node_comments, leading_alternate_branch_comments, leading_node_comments,
+ trailing_comments, trailing_node_comments,
+};
 use ruff_formatter::{SourceCode, SourceCodeSlice};
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::source_code::CommentRanges;
@@ -121,8 +124,6 @@ pub(crate) struct SourceComment {
  position: CommentTextPosition,
 }
 
-#[allow(unused)]
-// TODO(micha): Remove after using the new comments infrastructure in the formatter.
 impl SourceComment {
  /// Returns the location of the comment in the original source code.
  /// Allows retrieving the text of the comment.
@@ -184,8 +185,6 @@ pub(crate) enum CommentTextPosition {
  OwnLine,
 }
 
-#[allow(unused)]
-// TODO(micha): Remove after using the new comments infrastructure in the formatter.
 impl CommentTextPosition {
  pub(crate) const fn is_own_line(self) -> bool {
  matches!(self, CommentTextPosition::OwnLine)

crates/ruff_python_formatter/src/comments/placement.rs

@@ -5,7 +5,7 @@ use crate::trivia::find_first_non_trivia_character_in_range;
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::source_code::Locator;
 use ruff_python_ast::whitespace;
-use ruff_text_size::{TextRange, TextSize};
+use ruff_text_size::{TextLen, TextRange, TextSize};
 use rustpython_parser::ast::Ranged;
 use std::cmp::Ordering;
 
@@ -20,6 +20,7 @@ pub(super) fn place_comment<'a>(
  .or_else(|comment| handle_in_between_bodies_end_of_line_comment(comment, locator))
  .or_else(|comment| handle_trailing_body_comment(comment, locator))
  .or_else(handle_trailing_end_of_line_body_comment)
+ .or_else(|comment| handle_trailing_end_of_line_condition_comment(comment, locator))
  .or_else(|comment| handle_positional_only_arguments_separator_comment(comment, locator))
  .or_else(|comment| {
  handle_trailing_binary_expression_left_or_operator_comment(comment, locator)
@@ -471,6 +472,91 @@ fn handle_trailing_end_of_line_body_comment(comment: DecoratedComment<'_>) -> Co
  }
 }
 
+/// Handles end of line comments after the `:` of a condition
+///
+/// ```python
+/// while True: # comment
+/// pass
+/// ```
+///
+/// It attaches the comment as dangling comment to the enclosing `while` statement.
+fn handle_trailing_end_of_line_condition_comment<'a>(
+ comment: DecoratedComment<'a>,
+ locator: &Locator,
+) -> CommentPlacement<'a> {
+ use ruff_python_ast::prelude::*;
+
+ // Must be an end of line comment
+ if comment.text_position().is_own_line() {
+ return CommentPlacement::Default(comment);
+ }
+
+ // Must be between the condition expression and the first body element
+ let (Some(preceding), Some(following)) = (comment.preceding_node(), comment.following_node()) else {
+ return CommentPlacement::Default(comment);
+ };
+
+ let expression_before_colon = match comment.enclosing_node() {
+ AnyNodeRef::StmtIf(StmtIf { test: expr, .. })
+ | AnyNodeRef::StmtWhile(StmtWhile { test: expr, .. })
+ | AnyNodeRef::StmtFor(StmtFor { iter: expr, .. })
+ | AnyNodeRef::StmtAsyncFor(StmtAsyncFor { iter: expr, .. }) => {
+ Some(AnyNodeRef::from(expr.as_ref()))
+ }
+
+ AnyNodeRef::StmtWith(StmtWith { items, .. })
+ | AnyNodeRef::StmtAsyncWith(StmtAsyncWith { items, .. }) => {
+ items.last().map(AnyNodeRef::from)
+ }
+ _ => None,
+ };
+
+ let Some(last_before_colon) = expression_before_colon else {
+ return CommentPlacement::Default(comment);
+ };
+
+ // If the preceding is the node before the `colon`
+ // `while true:` The node before the `colon` is the `true` constant.
+ if preceding.ptr_eq(last_before_colon) {
+ let mut start = preceding.end();
+ while let Some((offset, c)) = find_first_non_trivia_character_in_range(
+ locator.contents(),
+ TextRange::new(start, following.start()),
+ ) {
+ match c {
+ ':' => {
+ if comment.slice().start() > offset {
+ // Comment comes after the colon
+ // ```python
+ // while a: # comment
+ // ...
+ // ```
+ return CommentPlacement::dangling(comment.enclosing_node(), comment);
+ }
+
+ // Comment comes before the colon
+ // ```python
+ // while (
+ // a # comment
+ // ):
+ // ...
+ // ```
+ break;
+ }
+ ')' => {
+ // Skip over any closing parentheses
+ start = offset + ')'.text_len();
+ }
+ _ => {
+ unreachable!("Only ')' or ':' should follow the condition")
+ }
+ }
+ }
+ }
+
+ CommentPlacement::Default(comment)
+}
+
 /// Attaches comments for the positional-only arguments separator `/` as trailing comments to the
 /// enclosing [`Arguments`] node.
 ///

crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__base_test.snap

@@ -19,19 +19,19 @@ expression: comments.debug(test_case.source_code)
  "trailing": [],
  },
  Node {
- kind: ExprCompare,
- range: 51..57,
- source: `x == y`,
+ kind: StmtIf,
+ range: 48..212,
+ source: `if x == y: # if statement e...ne comment⏎`,
  }: {
  "leading": [],
- "dangling": [],
- "trailing": [
+ "dangling": [
  SourceComment {
  text: "# if statement end of line comment",
  position: EndOfLine,
  formatted: false,
  },
  ],
+ "trailing": [],
  },
  Node {
  kind: StmtIf,