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

Simplify suite formatting (#6722)

Browse source 
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
1 changed files with 248 additions and 269 deletions
Download patch file Download diff file Expand all files Collapse all files

517
crates/ruff_python_formatter/src/statement/suite.rs
View file

@ -59,279 +59,258 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
let source = f.context().source(); let source = f.context().source();
let source_type = f.options().source_type(); let source_type = f.options().source_type();
let mut f = WithNodeLevel::new(node_level, f); let f = &mut WithNodeLevel::new(node_level, f);
write!(
f,
[format_with(|f| {
let mut iter = statements.iter();
let Some(first) = iter.next() else {
return Ok(());
};
// Format the first statement in the body, which often has special formatting rules. let mut iter = statements.iter();
let first = match self.kind { let Some(first) = iter.next() else {
SuiteKind::Other => { return Ok(());
if is_class_or_function_definition(first) && !comments.has_leading(first) { };
// Add an empty line for any nested functions or classes defined within
// non-function or class compound statements, e.g., this is stable formatting:
// ```python
// if True:
//
// def test():
// ...
// ```
empty_line().fmt(f)?;
}
SuiteChildStatement::Other(first) // Format the first statement in the body, which often has special formatting rules.
} let first = match self.kind {
SuiteKind::Other => {
SuiteKind::Function => { if is_class_or_function_definition(first) && !comments.has_leading(first) {
if let Some(docstring) = DocstringStmt::try_from_statement(first) { // Add an empty line for any nested functions or classes defined within
SuiteChildStatement::Docstring(docstring) // non-function or class compound statements, e.g., this is stable formatting:
} else { // ```python
SuiteChildStatement::Other(first) // if True:
} //
} // def test():
// ...
SuiteKind::Class => { // ```
if let Some(docstring) = DocstringStmt::try_from_statement(first) { empty_line().fmt(f)?;
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
// class Test:
//
// """Docstring"""
// ```
empty_line().fmt(f)?;
}
SuiteChildStatement::Docstring(docstring)
} else {
SuiteChildStatement::Other(first)
}
}
SuiteKind::TopLevel => SuiteChildStatement::Other(first),
};
let first_comments = comments.leading_dangling_trailing(first);
let (mut preceding, mut after_class_docstring) = if first_comments
.leading
.iter()
.any(|comment| comment.is_suppression_off_comment(source))
{
(
write_suppressed_statements_starting_with_leading_comment(
first, &mut iter, f,
)?,
false,
)
} else if first_comments
.trailing
.iter()
.any(|comment| comment.is_suppression_off_comment(source))
{
(
write_suppressed_statements_starting_with_trailing_comment(
first, &mut iter, f,
)?,
false,
)
} else {
first.fmt(f)?;
(
first.statement(),
matches!(first, SuiteChildStatement::Docstring(_))
&& matches!(self.kind, SuiteKind::Class),
)
};
while let Some(following) = iter.next() {
if is_class_or_function_definition(preceding)
|| is_class_or_function_definition(following)
{
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)
{
empty_line().fmt(f)?;
} else {
// Two subsequent classes that both have an ellipsis only body
// ```python
// class A: ...
// class B: ...
// ```
let class_sequences_with_ellipsis_only =
preceding.as_class_def_stmt().is_some_and(|class| {
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(),
)
});
// Two subsequent functions where the preceding has an ellipsis only body
// ```python
// def test(): ...
// def b(): a
// ```
let function_with_ellipsis =
preceding.as_function_def_stmt().is_some_and(|function| {
contains_only_an_ellipsis(
&function.body,
f.context().comments(),
)
}) && following.is_function_def_stmt();
// 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
{
empty_line().fmt(f)?;
}
}
}
SuiteKind::TopLevel => {
write!(f, [empty_line(), empty_line()])?;
}
SuiteKind::Function | SuiteKind::Class | SuiteKind::Other => {
empty_line().fmt(f)?;
}
}
} else if is_import_definition(preceding) && !is_import_definition(following) {
empty_line().fmt(f)?;
} else if is_compound_statement(preceding) {
// Handles the case where a body has trailing comments. The issue is that RustPython does not include
// the comments in the range of the suite. This means, the body ends right after the last statement in the body.
// ```python
// def test():
// ...
// # The body of `test` ends right after `...` and before this comment
//
// # leading comment
//
//
// a = 10
// ```
// Using `lines_after` for the node doesn't work because it would count the lines after the `...`
// 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()
{
first_leading.slice().start()
} else {
following.start()
};
match lines_before(start, source) {
0 | 1 => hard_line_break().fmt(f)?,
2 => empty_line().fmt(f)?,
3.. => match self.kind {
SuiteKind::TopLevel => write!(f, [empty_line(), empty_line()])?,
SuiteKind::Function | SuiteKind::Class | SuiteKind::Other => {
empty_line().fmt(f)?;
}
},
}
} else if after_class_docstring {
// Enforce an empty line after a class docstring, e.g., these are both stable
// formatting:
// ```python
// class Test:
// """Docstring"""
//
// ...
//
//
// class Test:
//
// """Docstring"""
//
// ...
// ```
empty_line().fmt(f)?;
after_class_docstring = false;
} else {
// Insert the appropriate number of empty lines based on the node level, e.g.:
// * [`NodeLevel::Module`]: Up to two empty lines
// * [`NodeLevel::CompoundStatement`]: Up to one empty line
// * [`NodeLevel::Expression`]: No empty lines
let count_lines = |offset| {
// It's necessary to skip any trailing line comment because RustPython doesn't include trailing comments
// in the node's range
// ```python
// a # The range of `a` ends right before this comment
//
// b
// ```
//
// Simply using `lines_after` doesn't work if a statement has a trailing comment because
// it then counts the lines between the statement and the trailing comment, which is
// always 0. This is why it skips any trailing trivia (trivia that's on the same line)
// and counts the lines after.
lines_after_ignoring_trivia(offset, source)
};
match node_level {
NodeLevel::TopLevel => match count_lines(preceding.end()) {
0 | 1 => hard_line_break().fmt(f)?,
2 => empty_line().fmt(f)?,
_ => write!(f, [empty_line(), empty_line()])?,
},
NodeLevel::CompoundStatement => match count_lines(preceding.end()) {
0 | 1 => hard_line_break().fmt(f)?,
_ => empty_line().fmt(f)?,
},
NodeLevel::Expression(_) | NodeLevel::ParenthesizedExpression => {
hard_line_break().fmt(f)?;
}
}
}
let following_comments = comments.leading_dangling_trailing(following);
if following_comments
.leading
.iter()
.any(|comment| comment.is_suppression_off_comment(source))
{
preceding = write_suppressed_statements_starting_with_leading_comment(
SuiteChildStatement::Other(following),
&mut iter,
f,
)?;
} else if following_comments
.trailing
.iter()
.any(|comment| comment.is_suppression_off_comment(source))
{
preceding = write_suppressed_statements_starting_with_trailing_comment(
SuiteChildStatement::Other(following),
&mut iter,
f,
)?;
} else {
following.format().fmt(f)?;
preceding = following;
}
} }
Ok(()) SuiteChildStatement::Other(first)
})] }
)
SuiteKind::Function => {
if let Some(docstring) = DocstringStmt::try_from_statement(first) {
SuiteChildStatement::Docstring(docstring)
} else {
SuiteChildStatement::Other(first)
}
}
SuiteKind::Class => {
if let Some(docstring) = DocstringStmt::try_from_statement(first) {
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
// class Test:
//
// """Docstring"""
// ```
empty_line().fmt(f)?;
}
SuiteChildStatement::Docstring(docstring)
} else {
SuiteChildStatement::Other(first)
}
}
SuiteKind::TopLevel => SuiteChildStatement::Other(first),
};
let first_comments = comments.leading_dangling_trailing(first);
let (mut preceding, mut after_class_docstring) = if first_comments
.leading
.iter()
.any(|comment| comment.is_suppression_off_comment(source))
{
(
write_suppressed_statements_starting_with_leading_comment(first, &mut iter, f)?,
false,
)
} else if first_comments
.trailing
.iter()
.any(|comment| comment.is_suppression_off_comment(source))
{
(
write_suppressed_statements_starting_with_trailing_comment(first, &mut iter, f)?,
false,
)
} else {
first.fmt(f)?;
(
first.statement(),
matches!(first, SuiteChildStatement::Docstring(_))
&& matches!(self.kind, SuiteKind::Class),
)
};
while let Some(following) = iter.next() {
if is_class_or_function_definition(preceding)
|| is_class_or_function_definition(following)
{
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) {
empty_line().fmt(f)?;
} else {
// Two subsequent classes that both have an ellipsis only body
// ```python
// class A: ...
// class B: ...
// ```
let class_sequences_with_ellipsis_only =
preceding.as_class_def_stmt().is_some_and(|class| {
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())
});
// Two subsequent functions where the preceding has an ellipsis only body
// ```python
// def test(): ...
// def b(): a
// ```
let function_with_ellipsis =
preceding.as_function_def_stmt().is_some_and(|function| {
contains_only_an_ellipsis(
&function.body,
f.context().comments(),
)
}) && following.is_function_def_stmt();
// 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 {
empty_line().fmt(f)?;
}
}
}
SuiteKind::TopLevel => {
write!(f, [empty_line(), empty_line()])?;
}
SuiteKind::Function | SuiteKind::Class | SuiteKind::Other => {
empty_line().fmt(f)?;
}
}
} else if is_import_definition(preceding) && !is_import_definition(following) {
empty_line().fmt(f)?;
} else if is_compound_statement(preceding) {
// Handles the case where a body has trailing comments. The issue is that RustPython does not include
// the comments in the range of the suite. This means, the body ends right after the last statement in the body.
// ```python
// def test():
// ...
// # The body of `test` ends right after `...` and before this comment
//
// # leading comment
//
//
// a = 10
// ```
// Using `lines_after` for the node doesn't work because it would count the lines after the `...`
// 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() {
first_leading.slice().start()
} else {
following.start()
};
match lines_before(start, source) {
0 | 1 => hard_line_break().fmt(f)?,
2 => empty_line().fmt(f)?,
3.. => match self.kind {
SuiteKind::TopLevel => write!(f, [empty_line(), empty_line()])?,
SuiteKind::Function | SuiteKind::Class | SuiteKind::Other => {
empty_line().fmt(f)?;
}
},
}
} else if after_class_docstring {
// Enforce an empty line after a class docstring, e.g., these are both stable
// formatting:
// ```python
// class Test:
// """Docstring"""
//
// ...
//
//
// class Test:
//
// """Docstring"""
//
// ...
// ```
empty_line().fmt(f)?;
after_class_docstring = false;
} else {
// Insert the appropriate number of empty lines based on the node level, e.g.:
// * [`NodeLevel::Module`]: Up to two empty lines
// * [`NodeLevel::CompoundStatement`]: Up to one empty line
// * [`NodeLevel::Expression`]: No empty lines
let count_lines = |offset| {
// It's necessary to skip any trailing line comment because RustPython doesn't include trailing comments
// in the node's range
// ```python
// a # The range of `a` ends right before this comment
//
// b
// ```
//
// Simply using `lines_after` doesn't work if a statement has a trailing comment because
// it then counts the lines between the statement and the trailing comment, which is
// always 0. This is why it skips any trailing trivia (trivia that's on the same line)
// and counts the lines after.
lines_after_ignoring_trivia(offset, source)
};
match node_level {
NodeLevel::TopLevel => match count_lines(preceding.end()) {
0 | 1 => hard_line_break().fmt(f)?,
2 => empty_line().fmt(f)?,
_ => write!(f, [empty_line(), empty_line()])?,
},
NodeLevel::CompoundStatement => match count_lines(preceding.end()) {
0 | 1 => hard_line_break().fmt(f)?,
_ => empty_line().fmt(f)?,
},
NodeLevel::Expression(_) | NodeLevel::ParenthesizedExpression => {
hard_line_break().fmt(f)?;
}
}
}
let following_comments = comments.leading_dangling_trailing(following);
if following_comments
.leading
.iter()
.any(|comment| comment.is_suppression_off_comment(source))
{
preceding = write_suppressed_statements_starting_with_leading_comment(
SuiteChildStatement::Other(following),
&mut iter,
f,
)?;
} else if following_comments
.trailing
.iter()
.any(|comment| comment.is_suppression_off_comment(source))
{
preceding = write_suppressed_statements_starting_with_trailing_comment(
SuiteChildStatement::Other(following),
&mut iter,
f,
)?;
} else {
following.format().fmt(f)?;
preceding = following;
}
}
Ok(())
} }
} }