mirror of
https://github.com/astral-sh/ruff.git
synced 2025-07-23 13:05:06 +00:00
formatter: WithNodeLevel helper (#6212)
This commit is contained in:
Micha Reiser
GitHub
parent
615337a54d
commit
38b5726948
8 changed files with 208 additions and 212 deletions
|
|
@ -2,7 +2,7 @@ use ruff_python_ast::{Expr, StmtAssign};
|
|||
|
||||
use ruff_formatter::{format_args, write, FormatError};
|
||||
|
||||
use crate::context::NodeLevel;
|
||||
use crate::context::{NodeLevel, WithNodeLevel};
|
||||
use crate::expression::parentheses::{Parentheses, Parenthesize};
|
||||
use crate::expression::{has_own_parentheses, maybe_parenthesize_expression};
|
||||
use crate::prelude::*;
|
||||
|
|
@ -61,13 +61,10 @@ impl Format<PyFormatContext<'_>> for FormatTargets<'_> {
|
|||
None
|
||||
};
|
||||
|
||||
let saved_level = f.context().node_level();
|
||||
f.context_mut()
|
||||
.set_node_level(NodeLevel::Expression(group_id));
|
||||
|
||||
let format_first = format_with(|f: &mut PyFormatter| {
|
||||
let result = if can_omit_parentheses {
|
||||
first.format().with_options(Parentheses::Never).fmt(f)
|
||||
let mut f = WithNodeLevel::new(NodeLevel::Expression(group_id), f);
|
||||
if can_omit_parentheses {
|
||||
write!(f, [first.format().with_options(Parentheses::Never)])
|
||||
} else {
|
||||
write!(
|
||||
f,
|
||||
|
|
@ -77,11 +74,7 @@ impl Format<PyFormatContext<'_>> for FormatTargets<'_> {
|
|||
if_group_breaks(&text(")"))
|
||||
]
|
||||
)
|
||||
};
|
||||
|
||||
f.context_mut().set_node_level(saved_level);
|
||||
|
||||
result
|
||||
}
|
||||
});
|
||||
|
||||
write!(
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ use ruff_python_ast::helpers::is_compound_statement;
|
|||
use ruff_python_ast::{Ranged, Stmt, Suite};
|
||||
use ruff_python_trivia::{lines_after, lines_before, skip_trailing_trivia};
|
||||
|
||||
use crate::context::NodeLevel;
|
||||
use crate::context::{NodeLevel, WithNodeLevel};
|
||||
use crate::prelude::*;
|
||||
|
||||
/// Level at which the [`Suite`] appears in the source code.
|
||||
|
|
@ -45,125 +45,111 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
|
|||
let comments = f.context().comments().clone();
|
||||
let source = f.context().source();
|
||||
|
||||
let saved_level = f.context().node_level();
|
||||
f.context_mut().set_node_level(node_level);
|
||||
let mut iter = statements.iter();
|
||||
let Some(first) = iter.next() else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
// Wrap the entire formatting operation in a `format_with` to ensure that we restore
|
||||
// context regardless of whether an error occurs.
|
||||
let formatted = format_with(|f| {
|
||||
let mut iter = statements.iter();
|
||||
let Some(first) = iter.next() else {
|
||||
return Ok(());
|
||||
};
|
||||
let mut f = WithNodeLevel::new(node_level, f);
|
||||
// First entry has never any separator, doesn't matter which one we take.
|
||||
write!(f, [first.format()])?;
|
||||
|
||||
// First entry has never any separator, doesn't matter which one we take.
|
||||
write!(f, [first.format()])?;
|
||||
let mut last = first;
|
||||
|
||||
let mut last = first;
|
||||
|
||||
for statement in iter {
|
||||
if is_class_or_function_definition(last)
|
||||
|| is_class_or_function_definition(statement)
|
||||
{
|
||||
match self.level {
|
||||
SuiteLevel::TopLevel => {
|
||||
write!(f, [empty_line(), empty_line(), statement.format()])?;
|
||||
}
|
||||
SuiteLevel::Nested => {
|
||||
write!(f, [empty_line(), statement.format()])?;
|
||||
}
|
||||
for statement in iter {
|
||||
if is_class_or_function_definition(last) || is_class_or_function_definition(statement) {
|
||||
match self.level {
|
||||
SuiteLevel::TopLevel => {
|
||||
write!(f, [empty_line(), empty_line(), statement.format()])?;
|
||||
}
|
||||
} else if is_import_definition(last) && !is_import_definition(statement) {
|
||||
write!(f, [empty_line(), statement.format()])?;
|
||||
} else if is_compound_statement(last) {
|
||||
// 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_comments(statement).first() {
|
||||
first_leading.slice().start()
|
||||
} else {
|
||||
statement.start()
|
||||
};
|
||||
|
||||
match lines_before(start, source) {
|
||||
0 | 1 => write!(f, [hard_line_break()])?,
|
||||
2 => write!(f, [empty_line()])?,
|
||||
3.. => {
|
||||
if self.level.is_nested() {
|
||||
write!(f, [empty_line()])?;
|
||||
} else {
|
||||
write!(f, [empty_line(), empty_line()])?;
|
||||
}
|
||||
}
|
||||
SuiteLevel::Nested => {
|
||||
write!(f, [empty_line(), statement.format()])?;
|
||||
}
|
||||
|
||||
write!(f, [statement.format()])?;
|
||||
} 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.
|
||||
let after_trailing_trivia = skip_trailing_trivia(offset, source);
|
||||
lines_after(after_trailing_trivia, source)
|
||||
}
|
||||
} else if is_import_definition(last) && !is_import_definition(statement) {
|
||||
write!(f, [empty_line(), statement.format()])?;
|
||||
} else if is_compound_statement(last) {
|
||||
// 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_comments(statement).first() {
|
||||
first_leading.slice().start()
|
||||
} else {
|
||||
statement.start()
|
||||
};
|
||||
|
||||
match node_level {
|
||||
NodeLevel::TopLevel => match count_lines(last.end()) {
|
||||
0 | 1 => write!(f, [hard_line_break()])?,
|
||||
2 => write!(f, [empty_line()])?,
|
||||
_ => write!(f, [empty_line(), empty_line()])?,
|
||||
},
|
||||
NodeLevel::CompoundStatement => match count_lines(last.end()) {
|
||||
0 | 1 => write!(f, [hard_line_break()])?,
|
||||
_ => write!(f, [empty_line()])?,
|
||||
},
|
||||
NodeLevel::Expression(_) | NodeLevel::ParenthesizedExpression => {
|
||||
write!(f, [hard_line_break()])?;
|
||||
match lines_before(start, source) {
|
||||
0 | 1 => write!(f, [hard_line_break()])?,
|
||||
2 => write!(f, [empty_line()])?,
|
||||
3.. => {
|
||||
if self.level.is_nested() {
|
||||
write!(f, [empty_line()])?;
|
||||
} else {
|
||||
write!(f, [empty_line(), empty_line()])?;
|
||||
}
|
||||
}
|
||||
|
||||
write!(f, [statement.format()])?;
|
||||
}
|
||||
|
||||
last = statement;
|
||||
write!(f, [statement.format()])?;
|
||||
} 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.
|
||||
let after_trailing_trivia = skip_trailing_trivia(offset, source);
|
||||
lines_after(after_trailing_trivia, source)
|
||||
};
|
||||
|
||||
match node_level {
|
||||
NodeLevel::TopLevel => match count_lines(last.end()) {
|
||||
0 | 1 => write!(f, [hard_line_break()])?,
|
||||
2 => write!(f, [empty_line()])?,
|
||||
_ => write!(f, [empty_line(), empty_line()])?,
|
||||
},
|
||||
NodeLevel::CompoundStatement => match count_lines(last.end()) {
|
||||
0 | 1 => write!(f, [hard_line_break()])?,
|
||||
_ => write!(f, [empty_line()])?,
|
||||
},
|
||||
NodeLevel::Expression(_) | NodeLevel::ParenthesizedExpression => {
|
||||
write!(f, [hard_line_break()])?;
|
||||
}
|
||||
}
|
||||
|
||||
write!(f, [statement.format()])?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
});
|
||||
last = statement;
|
||||
}
|
||||
|
||||
let result = formatted.fmt(f);
|
||||
|
||||
f.context_mut().set_node_level(saved_level);
|
||||
|
||||
result
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue