Format StmtFor (#5163)

<!--
Thank you for contributing to Ruff! To help us out with reviewing,
please consider the following:

- Does this pull request include a summary of the change? (See below.)
- Does this pull request include a descriptive title?
- Does this pull request include references to any relevant issues?
-->

## Summary

format StmtFor

still trying to learn how to help out with the formatter. trying
something slightly more advanced than [break](#5158)

mostly copied form StmtWhile

## Test Plan

snapshots
This commit is contained in:
David Szotten 2023-06-21 22:00:31 +01:00 committed by GitHub
parent e71f044f0d
commit 1eccbbb60e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 392 additions and 134 deletions

View file

@ -1,12 +1,91 @@
use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
use crate::comments::{leading_alternate_branch_comments, trailing_comments};
use crate::expression::expr_tuple::TupleParentheses;
use crate::expression::parentheses::Parenthesize;
use crate::prelude::*;
use crate::{FormatNodeRule, PyFormatter};
use ruff_formatter::{write, Buffer, FormatResult};
use rustpython_parser::ast::StmtFor;
use ruff_python_ast::node::AstNode;
use rustpython_parser::ast::{Expr, Ranged, Stmt, StmtFor};
#[derive(Debug)]
struct ExprTupleWithoutParentheses<'a>(&'a Expr);
impl Format<PyFormatContext<'_>> for ExprTupleWithoutParentheses<'_> {
fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
match self.0 {
Expr::Tuple(expr_tuple) => expr_tuple
.format()
.with_options(TupleParentheses::StripInsideForLoop)
.fmt(f),
other => other.format().with_options(Parenthesize::IfBreaks).fmt(f),
}
}
}
#[derive(Default)]
pub struct FormatStmtFor;
impl FormatNodeRule<StmtFor> for FormatStmtFor {
fn fmt_fields(&self, item: &StmtFor, f: &mut PyFormatter) -> FormatResult<()> {
write!(f, [not_yet_implemented(item)])
let StmtFor {
range: _,
target,
iter,
body,
orelse,
type_comment: _,
} = 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(iter.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("for"),
space(),
ExprTupleWithoutParentheses(target.as_ref()),
space(),
text("in"),
space(),
iter.format().with_options(Parenthesize::IfBreaks),
text(":"),
trailing_comments(trailing_condition_comments),
block_indent(&body.format())
]
)?;
if orelse.is_empty() {
debug_assert!(or_else_comments.is_empty());
} else {
// 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,
[
leading_alternate_branch_comments(leading, body.last()),
text("else:"),
trailing_comments(trailing),
block_indent(&orelse.format())
]
)?;
}
Ok(())
}
fn fmt_dangling_comments(&self, _node: &StmtFor, _f: &mut PyFormatter) -> FormatResult<()> {
// Handled in `fmt_fields`
Ok(())
}
}