mirror of
https://github.com/astral-sh/ruff.git
synced 2025-07-23 04:55:09 +00:00
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
GitHub
parent
e032fbd2e7
commit
b182368008
1 changed files with 248 additions and 269 deletions
|
|
@ -59,279 +59,258 @@ 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 mut iter = statements.iter();
|
||||
let Some(first) = iter.next() else {
|
||||
return Ok(());
|
||||
};
|
||||
let f = &mut WithNodeLevel::new(node_level, f);
|
||||
|
||||
// Format the first statement in the body, which often has special formatting rules.
|
||||
let first = match self.kind {
|
||||
SuiteKind::Other => {
|
||||
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)?;
|
||||
}
|
||||
let mut iter = statements.iter();
|
||||
let Some(first) = iter.next() else {
|
||||
return 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;
|
||||
}
|
||||
// Format the first statement in the body, which often has special formatting rules.
|
||||
let first = match self.kind {
|
||||
SuiteKind::Other => {
|
||||
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)?;
|
||||
}
|
||||
|
||||
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(())
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue