mirrors/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-07-25 05:53:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Refactor StmtIf: Formatter and Linter (#5459)

Browse source 
## Summary

Previously, `StmtIf` was defined recursively as
```rust
pub struct StmtIf {
    pub range: TextRange,
    pub test: Box<Expr>,
    pub body: Vec<Stmt>,
    pub orelse: Vec<Stmt>,
}
```
Every `elif` was represented as an `orelse` with a single `StmtIf`. This
means that this representation couldn't differentiate between
```python
if cond1:
    x = 1
else:
    if cond2:
        x = 2
```
and 
```python
if cond1:
    x = 1
elif cond2:
    x = 2
```
It also makes many checks harder than they need to be because we have to
recurse just to iterate over an entire if-elif-else and because we're
lacking nodes and ranges on the `elif` and `else` branches.

We change the representation to a flat

```rust
pub struct StmtIf {
    pub range: TextRange,
    pub test: Box<Expr>,
    pub body: Vec<Stmt>,
    pub elif_else_clauses: Vec<ElifElseClause>,
}

pub struct ElifElseClause {
    pub range: TextRange,
    pub test: Option<Expr>,
    pub body: Vec<Stmt>,
}
```
where `test: Some(_)` represents an `elif` and `test: None` an else.

This representation is different tradeoff, e.g. we need to allocate the
`Vec<ElifElseClause>`, the `elif`s are now different than the `if`s
(which matters in rules where want to check both `if`s and `elif`s) and
the type system doesn't guarantee that the `test: None` else is actually
last. We're also now a bit more inconsistent since all other `else`,
those from `for`, `while` and `try`, still don't have nodes. With the
new representation some things became easier, e.g. finding the `elif`
token (we can use the start of the `ElifElseClause`) and formatting
comments for if-elif-else (no more dangling comments splitting, we only
have to insert the dangling comment after the colon manually and set
`leading_alternate_branch_comments`, everything else is taken of by
having nodes for each branch and the usual placement.rs fixups).

## Merge Plan

This PR requires coordination between the parser repo and the main ruff
repo. I've split the ruff part, into two stacked PRs which have to be
merged together (only the second one fixes all tests), the first for the
formatter to be reviewed by @michareiser and the second for the linter
to be reviewed by @charliermarsh.

* MH: Review and merge
https://github.com/astral-sh/RustPython-Parser/pull/20
* MH: Review and merge or move later in stack
https://github.com/astral-sh/RustPython-Parser/pull/21
* MH: Review and approve
https://github.com/astral-sh/RustPython-Parser/pull/22
* MH: Review and approve formatter PR
https://github.com/astral-sh/ruff/pull/5459
* CM: Review and approve linter PR
https://github.com/astral-sh/ruff/pull/5460
* Merge linter PR in formatter PR, fix ecosystem checks (ecosystem
checks can't run on the formatter PR and won't run on the linter PR, so
we need to merge them first)
 * Merge https://github.com/astral-sh/RustPython-Parser/pull/22
 * Create tag in the parser, update linter+formatter PR
 * Merge linter+formatter PR https://github.com/astral-sh/ruff/pull/5459

---------

Co-authored-by: Micha Reiser <micha@reiser.io>
This commit is contained in:
konsti 2023-07-18 13:40:15 +02:00 • committed by GitHub
parent 167b9356fa
commit 730e6b2b4c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
82 changed files with 2333 additions and 2009 deletions
Download patch file Download diff file Expand all files Collapse all files

224
crates/ruff_python_formatter/src/comments/placement.rs
View file

@ -292,12 +292,26 @@ fn handle_in_between_bodies_own_line_comment<'a>(
// if x == y:
// pass
// # I'm a leading comment of the `elif` statement.
// elif:
// elif True:
// print("nooop")
// ```
if following.is_stmt_if() || following.is_except_handler() {
// The `elif` or except handlers have their own body to which we can attach the leading comment
if following.is_except_handler() {
// The except handlers have their own body to which we can attach the leading comment
CommentPlacement::leading(following, comment)
} else if let AnyNodeRef::StmtIf(stmt_if) = comment.enclosing_node() {
if let Some(clause) = stmt_if
.elif_else_clauses
.iter()
.find(|clause| are_same_optional(following, clause.test.as_ref()))
{
CommentPlacement::leading(clause.into(), comment)
} else {
// Since we know we're between bodies and we know that the following node is
// not the condition of any `elif`, we know the next node must be the `else`
let else_clause = stmt_if.elif_else_clauses.last().unwrap();
debug_assert!(else_clause.test.is_none());
CommentPlacement::leading(else_clause.into(), comment)
}
} else {
// There are no bodies for the "else" branch and other bodies that are represented as a `Vec<Stmt>`.
// This means, there's no good place to attach the comments to.
@ -356,42 +370,42 @@ fn handle_in_between_bodies_end_of_line_comment<'a>(
}
if locator.contains_line_break(TextRange::new(preceding.end(), comment.slice().start())) {
// The `elif` or except handlers have their own body to which we can attach the trailing comment
// The except handlers have their own body to which we can attach the trailing comment
// ```python
// if test:
// a
// elif c: # comment
// b
// try:
// f() # comment
// except RuntimeError:
// raise
// ```
if following.is_except_handler() {
return CommentPlacement::trailing(following, comment);
} else if following.is_stmt_if() {
// We have to exclude for following if statements that are not elif by checking the
// indentation
// ```python
// if True:
// pass
// else: # Comment
// if False:
// pass
// pass
// ```
let base_if_indent =
whitespace::indentation_at_offset(locator, following.range().start());
let maybe_elif_indent = whitespace::indentation_at_offset(
locator,
comment.enclosing_node().range().start(),
);
if base_if_indent == maybe_elif_indent {
return CommentPlacement::trailing(following, comment);
}
// Handle the `else` of an `if`. It is special because we don't have a test but unlike other
// `else` (e.g. for `while`), we have a dedicated node.
// ```python
// if x == y:
// pass
// elif x < y:
// pass
// else: # 12 trailing else condition
// pass
// ```
if let AnyNodeRef::StmtIf(stmt_if) = comment.enclosing_node() {
if let Some(else_clause) = stmt_if.elif_else_clauses.last() {
if else_clause.test.is_none()
&& following.ptr_eq(else_clause.body.first().unwrap().into())
{
return CommentPlacement::dangling(else_clause.into(), comment);
}
}
}
// There are no bodies for the "else" branch and other bodies that are represented as a `Vec<Stmt>`.
// This means, there's no good place to attach the comments to.
// Make this a dangling comments and manually format the comment in
// in the enclosing node's formatting logic. For `try`, it's the formatters responsibility
// to correctly identify the comments for the `finally` and `orelse` block by looking
// at the comment's range.
// There are no bodies for the "else" branch (only `Vec<Stmt>`) expect for StmtIf, so
// we make this a dangling comments of the node containing the alternate branch and
// manually format the comment in that node's formatting logic. For `try`, it's the
// formatters responsibility to correctly identify the comments for the `finally` and
// `orelse` block by looking at the comment's range.
//
// ```python
// while x == y:
@ -425,6 +439,64 @@ fn handle_in_between_bodies_end_of_line_comment<'a>(
}
}
/// Without the `StmtIf` special, this function would just be the following:
/// ```ignore
/// if let Some(preceding_node) = comment.preceding_node() {
/// Some((preceding_node, last_child_in_body(preceding_node)?))
/// } else {
/// None
/// }
/// ```
/// We handle two special cases here:
/// ```python
/// if True:
/// pass
/// # Comment between if and elif/else clause, needs to be manually attached to the `StmtIf`
/// else:
/// pass
/// # Comment after the `StmtIf`, needs to be manually attached to the ElifElseClause
/// ```
/// The problem is that `StmtIf` spans the whole range (there is no "inner if" node), so the first
/// comment doesn't see it as preceding node, and the second comment takes the entire `StmtIf` when
/// it should only take the `ElifElseClause`
fn find_preceding_and_handle_stmt_if_special_cases<'a>(
comment: &DecoratedComment<'a>,
) -> Option<(AnyNodeRef<'a>, AnyNodeRef<'a>)> {
if let (stmt_if @ AnyNodeRef::StmtIf(stmt_if_inner), Some(AnyNodeRef::ElifElseClause(..))) =
(comment.enclosing_node(), comment.following_node())
{
if let Some(preceding_node @ AnyNodeRef::ElifElseClause(..)) = comment.preceding_node() {
// We're already after and elif or else, defaults work
Some((preceding_node, last_child_in_body(preceding_node)?))
} else {
// Special case 1: The comment is between if body and an elif/else clause. We have
// to handle this separately since StmtIf spans the entire range, so it's not the
// preceding node
Some((
stmt_if,
AnyNodeRef::from(stmt_if_inner.body.last().unwrap()),
))
}
} else if let Some(preceding_node @ AnyNodeRef::StmtIf(stmt_if_inner)) =
comment.preceding_node()
{
if let Some(clause) = stmt_if_inner.elif_else_clauses.last() {
// Special case 2: We're after an if statement and need to narrow the preceding
// down to the elif/else clause
Some((clause.into(), last_child_in_body(clause.into())?))
} else {
// After an if without any elif/else, defaults work
Some((preceding_node, last_child_in_body(preceding_node)?))
}
} else if let Some(preceding_node) = comment.preceding_node() {
// The normal case
Some((preceding_node, last_child_in_body(preceding_node)?))
} else {
// Only do something if the preceding node has a body (has indented statements).
None
}
}
/// Handles trailing comments at the end of a body block (or any other block that is indented).
/// ```python
/// def test():
@ -442,12 +514,9 @@ fn handle_trailing_body_comment<'a>(
return CommentPlacement::Default(comment);
}
// Only do something if the preceding node has a body (has indented statements).
let Some(preceding_node) = comment.preceding_node() else {
return CommentPlacement::Default(comment);
};
let Some(last_child) = last_child_in_body(preceding_node) else {
let Some((preceding_node, last_child)) =
find_preceding_and_handle_stmt_if_special_cases(&comment)
else {
return CommentPlacement::Default(comment);
};
@ -566,6 +635,22 @@ fn handle_trailing_end_of_line_body_comment<'a>(
return CommentPlacement::Default(comment);
};
// Handle the StmtIf special case
// ```python
// if True:
// pass
// elif True:
// pass # 14 end-of-line trailing `pass` comment, set preceding to the ElifElseClause
// ```
let preceding = if let AnyNodeRef::StmtIf(stmt_if) = preceding {
stmt_if
.elif_else_clauses
.last()
.map_or(preceding, AnyNodeRef::from)
} else {
preceding
};
// Recursively get the last child of statements with a body.
let last_children = std::iter::successors(last_child_in_body(preceding), |parent| {
last_child_in_body(*parent)
@ -600,20 +685,40 @@ fn handle_trailing_end_of_line_condition_comment<'a>(
return CommentPlacement::Default(comment);
}
// We handle trailing else comments separately because we the preceding node is None for their
// case
// ```python
// if True:
// pass
// else: # 12 trailing else condition
// pass
// ```
if let AnyNodeRef::ElifElseClause(ast::ElifElseClause {
body, test: None, ..
}) = comment.enclosing_node()
{
if comment.start() < body.first().unwrap().start() {
return CommentPlacement::dangling(comment.enclosing_node(), 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() {
let enclosing_node = comment.enclosing_node();
let expression_before_colon = match enclosing_node {
AnyNodeRef::ElifElseClause(ast::ElifElseClause {
test: Some(expr), ..
}) => Some(AnyNodeRef::from(expr)),
AnyNodeRef::StmtIf(ast::StmtIf { test: expr, .. })
| AnyNodeRef::StmtWhile(ast::StmtWhile { test: expr, .. })
| AnyNodeRef::StmtFor(ast::StmtFor { iter: expr, .. })
| AnyNodeRef::StmtAsyncFor(ast::StmtAsyncFor { iter: expr, .. }) => {
Some(AnyNodeRef::from(expr.as_ref()))
}
AnyNodeRef::StmtWith(ast::StmtWith { items, .. })
| AnyNodeRef::StmtAsyncWith(ast::StmtAsyncWith { items, .. }) => {
items.last().map(AnyNodeRef::from)
@ -656,7 +761,7 @@ fn handle_trailing_end_of_line_condition_comment<'a>(
// while a: # comment
// ...
// ```
return CommentPlacement::dangling(comment.enclosing_node(), comment);
return CommentPlacement::dangling(enclosing_node, comment);
}
// Comment comes before the colon
@ -1439,10 +1544,15 @@ fn last_child_in_body(node: AnyNodeRef) -> Option<AnyNodeRef> {
| AnyNodeRef::MatchCase(ast::MatchCase { body, .. })
| AnyNodeRef::ExceptHandlerExceptHandler(ast::ExceptHandlerExceptHandler {
body, ..
}) => body,
})
| AnyNodeRef::ElifElseClause(ast::ElifElseClause { body, .. }) => body,
AnyNodeRef::StmtIf(ast::StmtIf {
body,
elif_else_clauses,
..
}) => elif_else_clauses.last().map_or(body, |clause| &clause.body),
AnyNodeRef::StmtIf(ast::StmtIf { body, orelse, .. })
| AnyNodeRef::StmtFor(ast::StmtFor { body, orelse, .. })
AnyNodeRef::StmtFor(ast::StmtFor { body, orelse, .. })
| AnyNodeRef::StmtAsyncFor(ast::StmtAsyncFor { body, orelse, .. })
| AnyNodeRef::StmtWhile(ast::StmtWhile { body, orelse, .. }) => {
if orelse.is_empty() {
@ -1453,7 +1563,7 @@ fn last_child_in_body(node: AnyNodeRef) -> Option<AnyNodeRef> {
}
AnyNodeRef::StmtMatch(ast::StmtMatch { cases, .. }) => {
return cases.last().map(AnyNodeRef::from)
return cases.last().map(AnyNodeRef::from);
}
AnyNodeRef::StmtTry(ast::StmtTry {
@ -1498,8 +1608,26 @@ fn is_first_statement_in_enclosing_alternate_body(
enclosing: AnyNodeRef,
) -> bool {
match enclosing {
AnyNodeRef::StmtIf(ast::StmtIf { orelse, .. })
| AnyNodeRef::StmtFor(ast::StmtFor { orelse, .. })
AnyNodeRef::StmtIf(ast::StmtIf {
elif_else_clauses, ..
}) => {
for clause in elif_else_clauses {
if let Some(test) = &clause.test {
// `elif`, the following node is the test
if following.ptr_eq(test.into()) {
return true;
}
} else {
// `else`, there is no test and the following node is the first entry in the
// body
if following.ptr_eq(clause.body.first().unwrap().into()) {
return true;
}
}
}
false
}
AnyNodeRef::StmtFor(ast::StmtFor { orelse, .. })
| AnyNodeRef::StmtAsyncFor(ast::StmtAsyncFor { orelse, .. })
| AnyNodeRef::StmtWhile(ast::StmtWhile { orelse, .. }) => {
are_same_optional(following, orelse.first())

4
crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__base_test.snap
View file

@ -34,8 +34,8 @@ expression: comments.debug(test_case.source_code)
"trailing": [],
},
Node {
kind: StmtIf,
range: 144..212,
kind: ElifElseClause,
range: 144..177,
source: `elif x < y:⏎`,
}: {
"leading": [

27
crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__if_elif_else_comments.snap
View file

@ -24,8 +24,8 @@ expression: comments.debug(test_case.source_code)
],
},
Node {
kind: StmtIf,
range: 104..192,
kind: ElifElseClause,
range: 104..124,
source: `elif x < y:⏎`,
}: {
"leading": [
@ -35,13 +35,7 @@ expression: comments.debug(test_case.source_code)
formatted: false,
},
],
"dangling": [
SourceComment {
text: "# Leading else comment",
position: OwnLine,
formatted: false,
},
],
"dangling": [],
"trailing": [],
},
Node {
@ -59,6 +53,21 @@ expression: comments.debug(test_case.source_code)
},
],
},
Node {
kind: ElifElseClause,
range: 178..192,
source: `else:⏎`,
}: {
"leading": [
SourceComment {
text: "# Leading else comment",
position: OwnLine,
formatted: false,
},
],
"dangling": [],
"trailing": [],
},
Node {
kind: StmtPass,
range: 188..192,

30
crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__if_elif_if_else_comments.snap
View file

@ -3,21 +3,6 @@ source: crates/ruff_python_formatter/src/comments/mod.rs
expression: comments.debug(test_case.source_code)
---
{
Node {
kind: StmtIf,
range: 21..128,
source: `elif x < y:⏎`,
}: {
"leading": [],
"dangling": [
SourceComment {
text: "# Leading else comment",
position: OwnLine,
formatted: false,
},
],
"trailing": [],
},
Node {
kind: StmtIf,
range: 37..60,
@ -33,4 +18,19 @@ expression: comments.debug(test_case.source_code)
},
],
},
Node {
kind: ElifElseClause,
range: 114..128,
source: `else:⏎`,
}: {
"leading": [
SourceComment {
text: "# Leading else comment",
position: OwnLine,
formatted: false,
},
],
"dangling": [],
"trailing": [],
},
}

11
crates/ruff_python_formatter/src/comments/visitor.rs
View file

@ -2,8 +2,8 @@ use std::iter::Peekable;
use ruff_text_size::{TextRange, TextSize};
use rustpython_parser::ast::{
Alias, Arg, ArgWithDefault, Arguments, Comprehension, Decorator, ExceptHandler, Expr, Keyword,
MatchCase, Mod, Pattern, Ranged, Stmt, WithItem,
Alias, Arg, ArgWithDefault, Arguments, Comprehension, Decorator, ElifElseClause, ExceptHandler,
Expr, Keyword, MatchCase, Mod, Pattern, Ranged, Stmt, WithItem,
};
use ruff_formatter::{SourceCode, SourceCodeSlice};
@ -284,6 +284,13 @@ impl<'ast> PreorderVisitor<'ast> for CommentsVisitor<'ast> {
}
self.finish_node(pattern);
}
fn visit_elif_else_clause(&mut self, elif_else_clause: &'ast ElifElseClause) {
if self.start_node(elif_else_clause).is_traverse() {
walk_elif_else_clause(self, elif_else_clause);
}
self.finish_node(elif_else_clause);
}
}
fn text_position(comment_range: TextRange, source_code: SourceCode) -> CommentLinePosition {

40
crates/ruff_python_formatter/src/generated.rs
View file

@ -617,6 +617,46 @@ impl<'ast> IntoFormat<PyFormatContext<'ast>> for ast::StmtIf {
}
}
impl FormatRule<ast::ElifElseClause, PyFormatContext<'_>>
for crate::statement::stmt_if::FormatElifElseClause
{
#[inline]
fn fmt(
&self,
node: &ast::ElifElseClause,
f: &mut Formatter<PyFormatContext<'_>>,
) -> FormatResult<()> {
FormatNodeRule::<ast::ElifElseClause>::fmt(self, node, f)
}
}
impl<'ast> AsFormat<PyFormatContext<'ast>> for ast::ElifElseClause {
type Format<'a> = FormatRefWithRule<
'a,
ast::ElifElseClause,
crate::statement::stmt_if::FormatElifElseClause,
PyFormatContext<'ast>,
>;
fn format(&self) -> Self::Format<'_> {
FormatRefWithRule::new(
self,
crate::statement::stmt_if::FormatElifElseClause::default(),
)
}
}
impl<'ast> IntoFormat<PyFormatContext<'ast>> for ast::ElifElseClause {
type Format = FormatOwnedWithRule<
ast::ElifElseClause,
crate::statement::stmt_if::FormatElifElseClause,
PyFormatContext<'ast>,
>;
fn into_format(self) -> Self::Format {
FormatOwnedWithRule::new(
self,
crate::statement::stmt_if::FormatElifElseClause::default(),
)
}
}
impl FormatRule<ast::StmtWith, PyFormatContext<'_>>
for crate::statement::stmt_with::FormatStmtWith
{

1
crates/ruff_python_formatter/src/statement/mod.rs
View file

@ -64,6 +64,7 @@ impl FormatRule<Stmt, PyFormatContext<'_>> for FormatStmt {
Stmt::Pass(x) => x.format().fmt(f),
Stmt::Break(x) => x.format().fmt(f),
Stmt::Continue(x) => x.format().fmt(f),
Stmt::TypeAlias(_) => todo!(),
}
}
}

1
crates/ruff_python_formatter/src/statement/stmt_class_def.rs
View file

@ -18,6 +18,7 @@ impl FormatNodeRule<StmtClassDef> for FormatStmtClassDef {
bases,
keywords,
body,
type_params: _,
decorator_list,
} = item;

181
crates/ruff_python_formatter/src/statement/stmt_if.rs
View file

@ -1,91 +1,43 @@
use crate::comments::{leading_alternate_branch_comments, trailing_comments, SourceComment};
use crate::comments::{leading_alternate_branch_comments, trailing_comments};
use crate::expression::maybe_parenthesize_expression;
use crate::expression::parentheses::Parenthesize;
use crate::prelude::*;
use crate::FormatNodeRule;
use ruff_formatter::{write, FormatError};
use rustpython_parser::ast::{Ranged, Stmt, StmtIf, Suite};
use ruff_formatter::write;
use ruff_python_ast::node::AnyNodeRef;
use rustpython_parser::ast::{ElifElseClause, StmtIf};
#[derive(Default)]
pub struct FormatStmtIf;
impl FormatNodeRule<StmtIf> for FormatStmtIf {
fn fmt_fields(&self, item: &StmtIf, f: &mut PyFormatter) -> FormatResult<()> {
let StmtIf {
range: _,
test,
body,
elif_else_clauses,
} = item;
let comments = f.context().comments().clone();
let trailing_colon_comment = comments.dangling_comments(item);
let mut current = IfOrElIf::If(item);
let mut else_comments: &[SourceComment];
let mut last_node_of_previous_body = None;
write!(
f,
[
text("if"),
space(),
maybe_parenthesize_expression(test, item, Parenthesize::IfBreaks),
text(":"),
trailing_comments(trailing_colon_comment),
block_indent(&body.format())
]
)?;
loop {
let current_statement = current.statement();
let StmtIf {
test, body, orelse, ..
} = current_statement;
let first_statement = body.first().ok_or(FormatError::SyntaxError)?;
let trailing = comments.dangling_comments(current_statement);
let trailing_if_comments_end = trailing
.partition_point(|comment| comment.slice().start() < first_statement.start());
let (if_trailing_comments, trailing_alternate_comments) =
trailing.split_at(trailing_if_comments_end);
if current.is_elif() {
let elif_leading = comments.leading_comments(current_statement);
// Manually format the leading comments because the formatting bypasses `NodeRule::fmt`
write!(
f,
[
leading_alternate_branch_comments(elif_leading, last_node_of_previous_body),
source_position(current_statement.start())
]
)?;
}
write!(
f,
[
text(current.keyword()),
space(),
maybe_parenthesize_expression(test, current_statement, Parenthesize::IfBreaks),
text(":"),
trailing_comments(if_trailing_comments),
block_indent(&body.format())
]
)?;
// RustPython models `elif` by setting the body to a single `if` statement. The `orelse`
// of the most inner `if` statement then becomes the `else` of the whole `if` chain.
// That's why it's necessary to take the comments here from the most inner `elif`.
else_comments = trailing_alternate_comments;
last_node_of_previous_body = body.last();
if let Some(elif) = else_if(orelse) {
current = elif;
} else {
break;
}
}
let orelse = &current.statement().orelse;
if !orelse.is_empty() {
// Leading comments are always own line comments
let leading_else_comments_end =
else_comments.partition_point(|comment| comment.line_position().is_own_line());
let (else_leading, else_trailing) = else_comments.split_at(leading_else_comments_end);
write!(
f,
[
leading_alternate_branch_comments(else_leading, last_node_of_previous_body),
text("else:"),
trailing_comments(else_trailing),
block_indent(&orelse.format())
]
)?;
let mut last_node = body.last().unwrap().into();
for clause in elif_else_clauses {
format_elif_else_clause(clause, f, Some(last_node))?;
last_node = clause.body.last().unwrap().into();
}
Ok(())
@ -97,35 +49,56 @@ impl FormatNodeRule<StmtIf> for FormatStmtIf {
}
}
fn else_if(or_else: &Suite) -> Option<IfOrElIf> {
if let [Stmt::If(if_stmt)] = or_else.as_slice() {
Some(IfOrElIf::ElIf(if_stmt))
/// Note that this implementation misses the leading newlines before the leading comments because
/// it does not have access to the last node of the previous branch. The `StmtIf` therefore doesn't
/// call this but `format_elif_else_clause` directly.
#[derive(Default)]
pub struct FormatElifElseClause;
impl FormatNodeRule<ElifElseClause> for FormatElifElseClause {
fn fmt_fields(&self, item: &ElifElseClause, f: &mut PyFormatter) -> FormatResult<()> {
format_elif_else_clause(item, f, None)
}
}
/// Extracted so we can implement `FormatElifElseClause` but also pass in `last_node` from
/// `FormatStmtIf`
fn format_elif_else_clause(
item: &ElifElseClause,
f: &mut PyFormatter,
last_node: Option<AnyNodeRef>,
) -> FormatResult<()> {
let ElifElseClause {
range: _,
test,
body,
} = item;
let comments = f.context().comments().clone();
let trailing_colon_comment = comments.dangling_comments(item);
let leading_comments = comments.leading_comments(item);
leading_alternate_branch_comments(leading_comments, last_node).fmt(f)?;
if let Some(test) = test {
write!(
f,
[
text("elif"),
space(),
maybe_parenthesize_expression(test, item, Parenthesize::IfBreaks),
]
)?;
} else {
None
}
}
enum IfOrElIf<'a> {
If(&'a StmtIf),
ElIf(&'a StmtIf),
}
impl<'a> IfOrElIf<'a> {
const fn statement(&self) -> &'a StmtIf {
match self {
IfOrElIf::If(statement) => statement,
IfOrElIf::ElIf(statement) => statement,
}
}
const fn keyword(&self) -> &'static str {
match self {
IfOrElIf::If(_) => "if",
IfOrElIf::ElIf(_) => "elif",
}
}
const fn is_elif(&self) -> bool {
matches!(self, IfOrElIf::ElIf(_))
text("else").fmt(f)?;
}
write!(
f,
[
text(":"),
trailing_comments(trailing_colon_comment),
block_indent(&body.format())
]
)
}