Format while Statement (#4810)

crates/ruff_python_formatter/src/statement/stmt_while.rs

@@ -1,12 +1,68 @@
-use crate::{verbatim_text, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
-use rustpython_parser::ast::StmtWhile;
+use crate::comments::{leading_alternate_branch_comments, trailing_comments};
+use crate::expression::maybe_parenthesize::maybe_parenthesize;
+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
     }
 }