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()) {
+            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 {
+    FormatLeadingComments::Node(node.as_any_node_ref())
-        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> {
+pub(crate) enum FormatLeadingComments<'a> {
-    node: AnyNodeRef<'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 trailing comments of `node`
+/// Formats the leading `comments` of an alternate branch and ensures that it preserves the right
-pub(crate) fn trailing_comments<T>(node: &T) -> FormatTrailingComments
+/// 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: AstNode,
+    T: Into<AnyNodeRef<'a>>,
 {
-    FormatTrailingComments {
+    FormatLeadingAlternateBranchComments {
-        node: node.as_any_node_ref(),
+        comments,
+        last_node: last_node.map(std::convert::Into::into),
     }
 }
 
-pub(crate) struct FormatTrailingComments<'a> {
+pub(crate) struct FormatLeadingAlternateBranchComments<'a> {
-    node: AnyNodeRef<'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_node_comments<T>(node: &T) -> FormatTrailingComments
+where
+    T: AstNode,
+{
+    FormatTrailingComments::Node(node.as_any_node_ref())
+}
+
+/// 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_trailing_own_line_comment |= trailing.position().is_own_line();
-            has_empty_lines_before |= lines_before_comment > 0;
+
+            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,
+        kind: StmtIf,
-        range: 51..57,
+        range: 48..212,
-        source: `x == y`,
+        source: `if x == y: # if statement e...ne comment⏎`,
     }: {
         "leading": [],
-        "dangling": [],
+        "dangling": [
-        "trailing": [
             SourceComment {
                 text: "# if statement end of line comment",
                 position: EndOfLine,
                 formatted: false,
             },
         ],
+        "trailing": [],
     },
     Node {
         kind: StmtIf,

crates/ruff_python_formatter/src/context.rs

@@ -25,7 +25,6 @@ impl<'a> PyFormatContext<'a> {
         }
     }
 
-    #[allow(unused)]
     pub(crate) fn contents(&self) -> &'a str {
         self.contents
     }
@@ -35,7 +34,6 @@ impl<'a> PyFormatContext<'a> {
         Locator::new(self.contents)
     }
 
-    #[allow(unused)]
     pub(crate) fn set_node_level(&mut self, level: NodeLevel) {
         self.node_level = level;
     }
@@ -44,7 +42,6 @@ impl<'a> PyFormatContext<'a> {
         self.node_level
     }
 
-    #[allow(unused)]
     pub(crate) fn comments(&self) -> &Comments<'a> {
         &self.comments
     }
@@ -80,11 +77,10 @@ pub(crate) enum NodeLevel {
     #[default]
     TopLevel,
 
-    /// Formatting nodes that are enclosed by a statement.
+    /// Formatting the body statements of a [compound statement](https://docs.python.org/3/reference/compound_stmts.html#compound-statements)
-    #[allow(unused)]
+    /// (`if`, `while`, `match`, etc.).
-    Statement,
+    CompoundStatement,
 
     /// Formatting nodes that are enclosed in a parenthesized expression.
-    #[allow(unused)]
     Parenthesized,
 }

crates/ruff_python_formatter/src/expression/maybe_parenthesize.rs

@@ -0,0 +1,53 @@
+use crate::context::NodeLevel;
+use crate::prelude::*;
+use ruff_formatter::{format_args, write};
+use rustpython_parser::ast::Expr;
+
+/// Formats the passed expression. Adds parentheses if the expression doesn't fit on a line.
+pub(crate) const fn maybe_parenthesize(expression: &Expr) -> MaybeParenthesize {
+    MaybeParenthesize { expression }
+}
+
+pub(crate) struct MaybeParenthesize<'a> {
+    expression: &'a Expr,
+}
+
+impl Format<PyFormatContext<'_>> for MaybeParenthesize<'_> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+        let saved_level = f.context().node_level();
+        f.context_mut().set_node_level(NodeLevel::Parenthesized);
+
+        let result = if needs_parentheses(self.expression) {
+            write!(
+                f,
+                [group(&format_args![
+                    if_group_breaks(&text("(")),
+                    soft_block_indent(&self.expression.format()),
+                    if_group_breaks(&text(")"))
+                ])]
+            )
+        } else {
+            // Don't add parentheses around expressions that  have parentheses on their own (e.g. list, dict, tuple, call expression)
+            self.expression.format().fmt(f)
+        };
+
+        f.context_mut().set_node_level(saved_level);
+
+        result
+    }
+}
+
+const fn needs_parentheses(expr: &Expr) -> bool {
+    !matches!(
+        expr,
+        Expr::Tuple(_)
+            | Expr::List(_)
+            | Expr::Set(_)
+            | Expr::Dict(_)
+            | Expr::ListComp(_)
+            | Expr::SetComp(_)
+            | Expr::DictComp(_)
+            | Expr::GeneratorExp(_)
+            | Expr::Call(_)
+    )
+}

crates/ruff_python_formatter/src/expression/mod.rs

@@ -29,6 +29,7 @@ pub(crate) mod expr_tuple;
 pub(crate) mod expr_unary_op;
 pub(crate) mod expr_yield;
 pub(crate) mod expr_yield_from;
+pub(crate) mod maybe_parenthesize;
 
 #[derive(Default)]
 pub struct FormatExpr;

crates/ruff_python_formatter/src/lib.rs

@@ -14,7 +14,9 @@ use ruff_formatter::{
 use ruff_python_ast::node::AstNode;
 use ruff_python_ast::source_code::{CommentRanges, CommentRangesBuilder, Locator};
 
-use crate::comments::{dangling_comments, leading_comments, trailing_comments, Comments};
+use crate::comments::{
+    dangling_node_comments, leading_node_comments, trailing_node_comments, Comments,
+};
 use crate::context::PyFormatContext;
 
 pub(crate) mod builders;
@@ -64,7 +66,7 @@ where
     /// You may want to override this method if you want to manually handle the formatting of comments
     /// inside of the `fmt_fields` method or customize the formatting of the leading comments.
     fn fmt_leading_comments(&self, node: &N, f: &mut PyFormatter) -> FormatResult<()> {
-        leading_comments(node).fmt(f)
+        leading_node_comments(node).fmt(f)
     }
 
     /// Formats the [dangling comments](comments#dangling-comments) of the node.
@@ -75,7 +77,7 @@ where
     ///
     /// A node can have dangling comments if all its children are tokens or if all node childrens are optional.
     fn fmt_dangling_comments(&self, node: &N, f: &mut PyFormatter) -> FormatResult<()> {
-        dangling_comments(node).fmt(f)
+        dangling_node_comments(node).fmt(f)
     }
 
     /// Formats the [trailing comments](comments#trailing-comments) of the node.
@@ -83,7 +85,7 @@ where
     /// You may want to override this method if you want to manually handle the formatting of comments
     /// inside of the `fmt_fields` method or customize the formatting of the trailing comments.
     fn fmt_trailing_comments(&self, node: &N, f: &mut PyFormatter) -> FormatResult<()> {
-        trailing_comments(node).fmt(f)
+        trailing_node_comments(node).fmt(f)
     }
 }
 
@@ -285,14 +287,58 @@ Formatted twice:
         Ok(())
     }
 
+    #[fixture(pattern = "resources/test/fixtures/ruff/**/*.py")]
+    #[test]
+    fn ruff_test(input_path: &Path) -> Result<()> {
+        let content = fs::read_to_string(input_path)?;
+
+        let printed = format_module(&content)?;
+        let formatted_code = printed.as_code();
+
+        let reformatted =
+            format_module(formatted_code).expect("Expected formatted code to be valid syntax");
+
+        if reformatted.as_code() != formatted_code {
+            let diff = TextDiff::from_lines(formatted_code, reformatted.as_code())
+                .unified_diff()
+                .header("Formatted once", "Formatted twice")
+                .to_string();
+            panic!(
+                r#"Reformatting the formatted code a second time resulted in formatting changes.
+{diff}
+
+Formatted once:
+{formatted_code}
+
+Formatted twice:
+{}"#,
+                reformatted.as_code()
+            );
+        }
+
+        let snapshot = format!(
+            r#"## Input
+{}
+
+## Output
+{}"#,
+            CodeFrame::new("py", &content),
+            CodeFrame::new("py", formatted_code)
+        );
+        assert_snapshot!(snapshot);
+
+        Ok(())
+    }
+
     /// Use this test to debug the formatting of some snipped
     #[ignore]
     #[test]
     fn quick_test() {
         let src = r#"
-{
+while True:
-    k: v for k, v in a_very_long_variable_name_that_exceeds_the_line_length_by_far_keep_going
+    if something.changed:
-}
+        do.stuff()  # trailing comment
+other
 "#;
         // Tokenize once
         let mut tokens = Vec::new();
@@ -320,10 +366,10 @@ Formatted twice:
 
         assert_eq!(
             printed.as_code(),
-            r#"{
+            r#"while True:
-    k: v
+    if something.changed:
-    for k, v in a_very_long_variable_name_that_exceeds_the_line_length_by_far_keep_going
+        do.stuff()  # trailing comment
-}"#
+"#
         );
     }
 

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments5_py.snap

@@ -86,20 +86,20 @@ if __name__ == "__main__":
 ```diff
 --- Black
 +++ Ruff
-@@ -1,11 +1,6 @@
+@@ -1,11 +1,9 @@
  while True:
      if something.changed:
 -        do.stuff()  # trailing comment
 -        # Comment belongs to the `if` block.
--    # This one belongs to the `while` block.
--
--    # Should this one, too?  I guess so.
--
 +        do.stuff()
+     # This one belongs to the `while` block.
+ 
+     # Should this one, too?  I guess so.
+-
  # This one is properly standalone now.
  
  for i in range(100):
-@@ -15,27 +10,18 @@
+@@ -15,27 +13,18 @@
  
      # then we do this
      print(i)
@@ -127,7 +127,7 @@ if __name__ == "__main__":
  # SECTION COMMENT
  
  
-@@ -47,8 +33,6 @@
+@@ -47,8 +36,6 @@
  @deco3
  def decorated1():
      ...
@@ -136,7 +136,7 @@ if __name__ == "__main__":
  # leading 1
  @deco1
  # leading 2
-@@ -56,18 +40,12 @@
+@@ -56,18 +43,12 @@
  # leading function comment
  def decorated1():
      ...
@@ -163,6 +163,9 @@ if __name__ == "__main__":
 while True:
     if something.changed:
         do.stuff()
+    # This one belongs to the `while` block.
+
+    # Should this one, too?  I guess so.
 # This one is properly standalone now.
 
 for i in range(100):

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__expression_py.snap

@@ -511,7 +511,7 @@ last_call()
  Ø = set()
  authors.łukasz.say_thanks()
  mapping = {
-@@ -233,138 +170,83 @@
+@@ -233,138 +170,84 @@
      C: 0.1 * (10.0 / 12),
      D: 0.1 * (10.0 / 12),
  }
@@ -550,15 +550,6 @@ last_call()
 -    ...
 -for j in 1 + (2 + 3):
 -    ...
--while this and that:
--    ...
--for (
--    addr_family,
--    addr_type,
--    addr_proto,
--    addr_canonname,
--    addr_sockaddr,
--) in socket.getaddrinfo("google.com", "http"):
 +print(* lambda x: x)
 +assert(not Test),("Short message")
 +assert this is ComplexTest and not requirements.fit_in_a_single_line(force=False), "Short message"
@@ -568,7 +559,15 @@ last_call()
 +for z in (i for i in (1, 2, 3)): ...
 +for i in (call()): ...
 +for j in (1 + (2 + 3)): ...
-+while(this and that): ...
+ while this and that:
+     ...
+-for (
+-    addr_family,
+-    addr_type,
+-    addr_proto,
+-    addr_canonname,
+-    addr_sockaddr,
+-) in socket.getaddrinfo("google.com", "http"):
 +for addr_family, addr_type, addr_proto, addr_canonname, addr_sockaddr in socket.getaddrinfo('google.com', 'http'):
      pass
 -a = (
@@ -879,7 +878,8 @@ for y in (): ...
 for z in (i for i in (1, 2, 3)): ...
 for i in (call()): ...
 for j in (1 + (2 + 3)): ...
-while(this and that): ...
+while this and that:
+    ...
 for addr_family, addr_type, addr_proto, addr_canonname, addr_sockaddr in socket.getaddrinfo('google.com', 'http'):
     pass
 a = aaaa.bbbb.cccc.dddd.eeee.ffff.gggg.hhhh.iiii.jjjj.kkkk.llll.mmmm.nnnn.oooo.pppp in qqqq.rrrr.ssss.tttt.uuuu.vvvv.xxxx.yyyy.zzzz

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__fmtskip8_py.snap

@@ -100,20 +100,24 @@ async def test_async_with():
  # Make sure a leading comment is not removed.
  if  unformatted_call(  args  ):  # fmt: skip
      print("First branch")
-@@ -30,33 +23,21 @@
+@@ -29,34 +22,22 @@
+ elif  another_unformatted_call(  args  ):  # fmt: skip
      print("Second branch")
  else  :  # fmt: skip
-     print("Last branch")
+-    print("Last branch")
 -
 -
- while  some_condition(  unformatted,  args  ):  # fmt: skip
+-while  some_condition(  unformatted,  args  ):  # fmt: skip
++    print("Last branch")  # fmt: skip
++while some_condition(  unformatted,  args  ):  # fmt: skip
      print("Do something")
 -
 -
  for  i  in  some_iter(  unformatted,  args  ):  # fmt: skip
-     print("Do something")
+-    print("Do something")
 -
 -
++    print("Do something")  # fmt: skip
  async def test_async_for():
      async  for  i  in  some_async_iter(  unformatted,  args  ):  # fmt: skip
          print("Do something")
@@ -128,9 +132,10 @@ async def test_async_with():
 -
 -
  with  give_me_context(  unformatted,  args  ):  # fmt: skip
-     print("Do something")
+-    print("Do something")
 -
 -
++    print("Do something")  # fmt: skip
  async def test_async_with():
      async  with  give_me_async_context(  unformatted,  args  ):  # fmt: skip
          print("Do something")
@@ -163,11 +168,11 @@ if  unformatted_call(  args  ):  # fmt: skip
 elif  another_unformatted_call(  args  ):  # fmt: skip
     print("Second branch")
 else  :  # fmt: skip
-    print("Last branch")
+    print("Last branch")  # fmt: skip
 while some_condition(  unformatted,  args  ):  # fmt: skip
     print("Do something")
 for  i  in  some_iter(  unformatted,  args  ):  # fmt: skip
-    print("Do something")
+    print("Do something")  # fmt: skip
 async def test_async_for():
     async  for  i  in  some_async_iter(  unformatted,  args  ):  # fmt: skip
         print("Do something")
@@ -178,7 +183,7 @@ except  UnformattedError  as  ex:  # fmt: skip
 finally  :  # fmt: skip
     finally_call()
 with  give_me_context(  unformatted,  args  ):  # fmt: skip
-    print("Do something")
+    print("Do something")  # fmt: skip
 async def test_async_with():
     async  with  give_me_async_context(  unformatted,  args  ):  # fmt: skip
         print("Do something")

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__remove_newline_after_code_block_open_py.snap

@@ -121,30 +121,30 @@ with open("/path/to/file.txt", mode="r") as read_file:
 ```diff
 --- Black
 +++ Ruff
-@@ -1,78 +1,74 @@
+@@ -1,78 +1,68 @@
  import random
--
++def foo1():
--
- def foo1():
--    print("The newline above me should be deleted!")
--
  
 +    print("The newline above me should be deleted!")
- def foo2():
++def foo2():
--    print("All the newlines above me should be deleted!")
+ 
+-def foo1():
+-    print("The newline above me should be deleted!")
  
  
-+
+-def foo2():
-+    print("All the newlines above me should be deleted!")
+     print("All the newlines above me should be deleted!")
+-
+-
  def foo3():
 +
      print("No newline above me!")
  
      print("There is a newline above me, and that's OK!")
-+def foo4():
- 
 -
--def foo4():
+-
+ def foo4():
++
      # There is a comment here
  
      print("The newline above me should not be deleted!")
@@ -154,23 +154,23 @@ with open("/path/to/file.txt", mode="r") as read_file:
      def bar(self):
 +
          print("The newline above me should be deleted!")
--
++for i in range(5):
--
- for i in range(5):
--    print(f"{i}) The line above me should be removed!")
--
  
 +    print(f"{i}) The line above me should be removed!")
++for i in range(5):
+ 
+-for i in range(5):
+-    print(f"{i}) The line above me should be removed!")
+ 
+ 
++    print(f"{i}) The lines above me should be removed!")
  for i in range(5):
 -    print(f"{i}) The lines above me should be removed!")
  
++    for j in range(7):
  
-+
+-for i in range(5):
-+    print(f"{i}) The lines above me should be removed!")
+-    for j in range(7):
- for i in range(5):
-+
-     for j in range(7):
-+
          print(f"{i}) The lines above me should be removed!")
 +if random.randint(0, 3) == 0:
  
@@ -189,38 +189,33 @@ with open("/path/to/file.txt", mode="r") as read_file:
      if random.uniform(0, 1) > 0.5:
          print("Two lines above me are about to be removed!")
 -
-+while True:
+-
- 
+ while True:
-+    print("The newline above me should be deleted!")
+     print("The newline above me should be deleted!")
+-
+-
  while True:
--    print("The newline above me should be deleted!")
- 
- 
--while True:
-+
      print("The newlines above me should be deleted!")
-+while True:
- 
 -
--while True:
+-
+ while True:
      while False:
--        print("The newlines above me should be deleted!")
+         print("The newlines above me should be deleted!")
--
++with open("/path/to/file.txt", mode="w") as file:
  
-+        print("The newlines above me should be deleted!")
- with open("/path/to/file.txt", mode="w") as file:
--    file.write("The new line above me is about to be removed!")
 -
- 
 +    file.write("The new line above me is about to be removed!")
  with open("/path/to/file.txt", mode="w") as file:
--    file.write("The new lines above me is about to be removed!")
+-    file.write("The new line above me is about to be removed!")
  
  
+-with open("/path/to/file.txt", mode="w") as file:
 +
-+    file.write("The new lines above me is about to be removed!")
+     file.write("The new lines above me is about to be removed!")
- with open("/path/to/file.txt", mode="r") as read_file:
++with open("/path/to/file.txt", mode="r") as read_file:
-+
+ 
+-
+-with open("/path/to/file.txt", mode="r") as read_file:
      with open("/path/to/output_file.txt", mode="w") as write_file:
 +
          write_file.writelines(read_file.readlines())
@@ -278,17 +273,11 @@ if random.randint(0, 3) == 0:
     if random.uniform(0, 1) > 0.5:
         print("Two lines above me are about to be removed!")
 while True:
-
     print("The newline above me should be deleted!")
 while True:
-
-
-
     print("The newlines above me should be deleted!")
 while True:
-
     while False:
-
         print("The newlines above me should be deleted!")
 with open("/path/to/file.txt", mode="w") as file:
 

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__ruff_test__while_py.snap

@@ -0,0 +1,70 @@
+---
+source: crates/ruff_python_formatter/src/lib.rs
+expression: snapshot
+---
+## Input
+```py
+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")
+```
+
+
+
+## Output
+```py
+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/statement/stmt_while.rs

@@ -1,12 +1,68 @@
-use crate::{verbatim_text, FormatNodeRule, PyFormatter};
+use crate::comments::{leading_alternate_branch_comments, trailing_comments};
-use ruff_formatter::{write, Buffer, FormatResult};
+use crate::expression::maybe_parenthesize::maybe_parenthesize;
-use rustpython_parser::ast::StmtWhile;
+use crate::prelude::*;
+use crate::FormatNodeRule;
+use ruff_formatter::write;
+use ruff_python_ast::node::AstNode;
+use rustpython_parser::ast::{Ranged, Stmt, StmtWhile};
 
 #[derive(Default)]
 pub struct FormatStmtWhile;
 
 impl FormatNodeRule<StmtWhile> for FormatStmtWhile {
     fn fmt_fields(&self, item: &StmtWhile, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [verbatim_text(item.range)])
+        let StmtWhile {
+            range: _,
+            test,
+            body,
+            orelse,
+        } = item;
+
+        let comments = f.context().comments().clone();
+        let dangling_comments = comments.dangling_comments(item.as_any_node_ref());
+
+        let body_start = body.first().map_or(test.end(), Stmt::start);
+        let or_else_comments_start =
+            dangling_comments.partition_point(|comment| comment.slice().end() < body_start);
+
+        let (trailing_condition_comments, or_else_comments) =
+            dangling_comments.split_at(or_else_comments_start);
+
+        write!(
+            f,
+            [
+                text("while"),
+                space(),
+                maybe_parenthesize(test),
+                text(":"),
+                trailing_comments(trailing_condition_comments),
+                block_indent(&body.format())
+            ]
+        )?;
+
+        if !orelse.is_empty() {
+            // Split between leading comments before the `else` keyword and end of line comments at the end of
+            // the `else:` line.
+            let trailing_start =
+                or_else_comments.partition_point(|comment| comment.position().is_own_line());
+            let (leading, trailing) = or_else_comments.split_at(trailing_start);
+
+            write!(
+                f,
+                [
+                    leading_alternate_branch_comments(leading, body.last()),
+                    text("else:"),
+                    trailing_comments(trailing),
+                    block_indent(&orelse.format())
+                ]
+            )?;
+        }
+
+        Ok(())
+    }
+
+    fn fmt_dangling_comments(&self, _node: &StmtWhile, _f: &mut PyFormatter) -> FormatResult<()> {
+        // Handled in `fmt_fields`
+        Ok(())
     }
 }

crates/ruff_python_formatter/src/statement/suite.rs

@@ -28,10 +28,15 @@ impl Default for FormatSuite {
 
 impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
     fn fmt(&self, statements: &Suite, f: &mut PyFormatter) -> FormatResult<()> {
-        let mut joiner = f.join_nodes(match self.level {
+        let node_level = match self.level {
             SuiteLevel::TopLevel => NodeLevel::TopLevel,
-            SuiteLevel::Nested => NodeLevel::Statement,
+            SuiteLevel::Nested => NodeLevel::CompoundStatement,
-        });
+        };
+
+        let saved_level = f.context().node_level();
+        f.context_mut().set_node_level(node_level);
+
+        let mut joiner = f.join_nodes(node_level);
 
         let mut iter = statements.iter();
         let Some(first) = iter.next() else {
@@ -67,7 +72,11 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
             is_last_function_or_class_definition = is_current_function_or_class_definition;
         }
 
-        joiner.finish()
+        let result = joiner.finish();
+
+        f.context_mut().set_node_level(saved_level);
+
+        result
     }
 }
 

crates/ruff_python_formatter/src/trivia.rs

@@ -41,7 +41,6 @@ pub(crate) fn find_first_non_trivia_character_in_range(
 }
 
 /// Returns the number of newlines between `offset` and the first non whitespace character in the source code.
-#[allow(unused)] // TODO(micha) Remove after using for statements.
 pub(crate) fn lines_before(code: &str, offset: TextSize) -> u32 {
     let head = &code[TextRange::up_to(offset)];
     let mut newlines = 0u32;