ruff/crates/ruff_python_formatter/src/statement/stmt_while.rs

69 lines
2.2 KiB
Rust

use crate::comments::{leading_alternate_branch_comments, trailing_comments};
use crate::expression::parentheses::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<()> {
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(),
test.format().with_options(Parenthesize::IfBreaks),
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(())
}
}