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(())
     }
 }