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,