ruff/crates/ruff_python_formatter/src/statement/stmt_while.rs
konsti 9a817a2922
Insert empty line between suite and alternative branch after def/class (#12294)
When there is a function or class definition at the end of a suite
followed by the beginning of an alternative block, we have to insert a
single empty line between them.

In the if-else-statement example below, we insert an empty line after
the `foo` in the if-block, but none after the else-block `foo`, since in
the latter case the enclosing suite already adds empty lines.

```python
if sys.version_info >= (3, 10):
    def foo():
        return "new"
else:
    def foo():
        return "old"
class Bar:
    pass
```

To do so, we track whether the current suite is the last one in the
current statement with a new option on the suite kind.

Fixes #12199

---------

Co-authored-by: Micha Reiser <micha@reiser.io>
2024-07-15 12:59:33 +02:00

77 lines
2.6 KiB
Rust

use ruff_formatter::{format_args, write};
use ruff_python_ast::AstNode;
use ruff_python_ast::{Stmt, StmtWhile};
use ruff_text_size::Ranged;
use crate::expression::maybe_parenthesize_expression;
use crate::expression::parentheses::Parenthesize;
use crate::prelude::*;
use crate::statement::clause::{clause_body, clause_header, ClauseHeader, ElseClause};
use crate::statement::suite::SuiteKind;
#[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(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.end() < body_start);
let (trailing_condition_comments, or_else_comments) =
dangling_comments.split_at(or_else_comments_start);
write!(
f,
[
clause_header(
ClauseHeader::While(item),
trailing_condition_comments,
&format_args![
token("while"),
space(),
maybe_parenthesize_expression(test, item, Parenthesize::IfBreaks),
]
),
clause_body(
body,
SuiteKind::other(orelse.is_empty()),
trailing_condition_comments
),
]
)?;
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.line_position().is_own_line());
let (leading, trailing) = or_else_comments.split_at(trailing_start);
write!(
f,
[
clause_header(
ClauseHeader::OrElse(ElseClause::While(item)),
trailing,
&token("else")
)
.with_leading_comments(leading, body.last()),
clause_body(orelse, SuiteKind::other(true), trailing),
]
)?;
}
Ok(())
}
}