Simplify suite formatting (#6722)

Avoid the nesting in a macro by using the new `WithNodeLevel` to
`PyFormatter` deref. No changes otherwise.

I wanted to follow this up with quickly fixing the typeshed empty line
rules but they turned out a lot more complex than i had anticipated.
This commit is contained in:
konsti 2023-08-21 21:01:51 +02:00 committed by GitHub
parent e032fbd2e7
commit b182368008
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -59,10 +59,8 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
let source = f.context().source();
let source_type = f.options().source_type();
let mut f = WithNodeLevel::new(node_level, f);
write!(
f,
[format_with(|f| {
let f = &mut WithNodeLevel::new(node_level, f);
let mut iter = statements.iter();
let Some(first) = iter.next() else {
return Ok(());
@ -96,9 +94,7 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
SuiteKind::Class => {
if let Some(docstring) = DocstringStmt::try_from_statement(first) {
if !comments.has_leading(first)
&& lines_before(first.start(), source) > 1
{
if !comments.has_leading(first) && lines_before(first.start(), source) > 1 {
// Allow up to one empty line before a class docstring, e.g., this is
// stable formatting:
// ```python
@ -125,9 +121,7 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
.any(|comment| comment.is_suppression_off_comment(source))
{
(
write_suppressed_statements_starting_with_leading_comment(
first, &mut iter, f,
)?,
write_suppressed_statements_starting_with_leading_comment(first, &mut iter, f)?,
false,
)
} else if first_comments
@ -136,9 +130,7 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
.any(|comment| comment.is_suppression_off_comment(source))
{
(
write_suppressed_statements_starting_with_trailing_comment(
first, &mut iter, f,
)?,
write_suppressed_statements_starting_with_trailing_comment(first, &mut iter, f)?,
false,
)
} else {
@ -157,9 +149,7 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
match self.kind {
SuiteKind::TopLevel if source_type.is_stub() => {
// Preserve the empty line if the definitions are separated by a comment
if comments.has_trailing(preceding)
|| comments.has_leading(following)
{
if comments.has_trailing(preceding) || comments.has_leading(following) {
empty_line().fmt(f)?;
} else {
// Two subsequent classes that both have an ellipsis only body
@ -169,15 +159,9 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
// ```
let class_sequences_with_ellipsis_only =
preceding.as_class_def_stmt().is_some_and(|class| {
contains_only_an_ellipsis(
&class.body,
f.context().comments(),
)
contains_only_an_ellipsis(&class.body, f.context().comments())
}) && following.as_class_def_stmt().is_some_and(|class| {
contains_only_an_ellipsis(
&class.body,
f.context().comments(),
)
contains_only_an_ellipsis(&class.body, f.context().comments())
});
// Two subsequent functions where the preceding has an ellipsis only body
@ -195,9 +179,7 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
// Don't add an empty line between two classes that have an `...` body only or after
// a function with an `...` body. Otherwise add an empty line.
if !class_sequences_with_ellipsis_only
&& !function_with_ellipsis
{
if !class_sequences_with_ellipsis_only && !function_with_ellipsis {
empty_line().fmt(f)?;
}
}
@ -228,8 +210,7 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
// which is 0 instead of 1, the number of lines between the trailing comment and
// the leading comment. This is why the suite handling counts the lines before the
// start of the next statement or before the first leading comments for compound statements.
let start = if let Some(first_leading) = comments.leading(following).first()
{
let start = if let Some(first_leading) = comments.leading(following).first() {
first_leading.slice().start()
} else {
following.start()
@ -330,8 +311,6 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
}
Ok(())
})]
)
}
}