mirror of
https://github.com/astral-sh/ruff.git
synced 2025-08-21 19:04:33 +00:00
Support fmt: skip on compound statements (#6593)
This commit is contained in:
Micha Reiser
GitHub
parent
4dc32a00d0
commit
fa7442da2f
23 changed files with 1385 additions and 625 deletions
407
crates/ruff_python_formatter/src/statement/clause.rs
Normal file
407
crates/ruff_python_formatter/src/statement/clause.rs
Normal file
|
|
@ -0,0 +1,407 @@
|
|||
use crate::comments::{
|
||||
leading_alternate_branch_comments, trailing_comments, SourceComment, SuppressionKind,
|
||||
};
|
||||
use crate::prelude::*;
|
||||
use crate::verbatim::write_suppressed_clause_header;
|
||||
use ruff_formatter::{Argument, Arguments, FormatError};
|
||||
use ruff_python_ast::node::AnyNodeRef;
|
||||
use ruff_python_ast::{
|
||||
ElifElseClause, ExceptHandlerExceptHandler, MatchCase, Ranged, StmtClassDef, StmtFor,
|
||||
StmtFunctionDef, StmtIf, StmtMatch, StmtTry, StmtWhile, StmtWith,
|
||||
};
|
||||
use ruff_python_trivia::{SimpleToken, SimpleTokenKind, SimpleTokenizer};
|
||||
use ruff_text_size::{TextRange, TextSize};
|
||||
|
||||
/// The header of a compound statement clause.
|
||||
///
|
||||
/// > A compound statement consists of one or more ‘clauses.’ A clause consists of a header and a ‘suite.’
|
||||
/// > The clause headers of a particular compound statement are all at the same indentation level.
|
||||
/// > Each clause header begins with a uniquely identifying keyword and ends with a colon.
|
||||
/// [source](https://docs.python.org/3/reference/compound_stmts.html#compound-statements)
|
||||
#[derive(Copy, Clone)]
|
||||
pub(crate) enum ClauseHeader<'a> {
|
||||
Class(&'a StmtClassDef),
|
||||
Function(&'a StmtFunctionDef),
|
||||
If(&'a StmtIf),
|
||||
ElifElse(&'a ElifElseClause),
|
||||
Try(&'a StmtTry),
|
||||
ExceptHandler(&'a ExceptHandlerExceptHandler),
|
||||
TryFinally(&'a StmtTry),
|
||||
Match(&'a StmtMatch),
|
||||
MatchCase(&'a MatchCase),
|
||||
For(&'a StmtFor),
|
||||
While(&'a StmtWhile),
|
||||
With(&'a StmtWith),
|
||||
OrElse(ElseClause<'a>),
|
||||
}
|
||||
|
||||
impl<'a> ClauseHeader<'a> {
|
||||
/// The range from the clause keyword up to and including the final colon.
|
||||
pub(crate) fn range(self, source: &str) -> FormatResult<TextRange> {
|
||||
let keyword_range = self.first_keyword_range(source)?;
|
||||
|
||||
let mut last_child_end = None;
|
||||
|
||||
self.visit(&mut |child| last_child_end = Some(child.end()));
|
||||
|
||||
let end = match self {
|
||||
ClauseHeader::Class(class) => Some(last_child_end.unwrap_or(class.name.end())),
|
||||
ClauseHeader::Function(function) => Some(last_child_end.unwrap_or(function.name.end())),
|
||||
ClauseHeader::ElifElse(_)
|
||||
| ClauseHeader::Try(_)
|
||||
| ClauseHeader::If(_)
|
||||
| ClauseHeader::TryFinally(_)
|
||||
| ClauseHeader::Match(_)
|
||||
| ClauseHeader::MatchCase(_)
|
||||
| ClauseHeader::For(_)
|
||||
| ClauseHeader::While(_)
|
||||
| ClauseHeader::With(_)
|
||||
| ClauseHeader::OrElse(_) => last_child_end,
|
||||
|
||||
ClauseHeader::ExceptHandler(handler) => handler
|
||||
.name
|
||||
.as_ref()
|
||||
.map(ruff_python_ast::Ranged::end)
|
||||
.or(last_child_end),
|
||||
};
|
||||
|
||||
let colon = colon_range(end.unwrap_or(keyword_range.end()), source)?;
|
||||
|
||||
Ok(TextRange::new(keyword_range.start(), colon.end()))
|
||||
}
|
||||
|
||||
/// Visits the nodes in the case header.
|
||||
pub(crate) fn visit<F>(self, visitor: &mut F)
|
||||
where
|
||||
F: FnMut(AnyNodeRef),
|
||||
{
|
||||
fn visit<'a, N, F>(node: N, visitor: &mut F)
|
||||
where
|
||||
N: Into<AnyNodeRef<'a>>,
|
||||
F: FnMut(AnyNodeRef<'a>),
|
||||
{
|
||||
visitor(node.into());
|
||||
}
|
||||
|
||||
match self {
|
||||
ClauseHeader::Class(StmtClassDef {
|
||||
type_params,
|
||||
arguments,
|
||||
range: _,
|
||||
decorator_list: _,
|
||||
name: _,
|
||||
body: _,
|
||||
}) => {
|
||||
if let Some(type_params) = type_params.as_deref() {
|
||||
visit(type_params, visitor);
|
||||
}
|
||||
|
||||
if let Some(arguments) = arguments {
|
||||
visit(arguments.as_ref(), visitor);
|
||||
}
|
||||
}
|
||||
ClauseHeader::Function(StmtFunctionDef {
|
||||
type_params,
|
||||
parameters,
|
||||
range: _,
|
||||
is_async: _,
|
||||
decorator_list: _,
|
||||
name: _,
|
||||
returns: _,
|
||||
body: _,
|
||||
}) => {
|
||||
if let Some(type_params) = type_params.as_ref() {
|
||||
visit(type_params, visitor);
|
||||
}
|
||||
visit(parameters.as_ref(), visitor);
|
||||
}
|
||||
ClauseHeader::If(StmtIf {
|
||||
test,
|
||||
range: _,
|
||||
body: _,
|
||||
elif_else_clauses: _,
|
||||
}) => {
|
||||
visit(test.as_ref(), visitor);
|
||||
}
|
||||
ClauseHeader::ElifElse(ElifElseClause {
|
||||
test,
|
||||
range: _,
|
||||
body: _,
|
||||
}) => {
|
||||
if let Some(test) = test.as_ref() {
|
||||
visit(test, visitor);
|
||||
}
|
||||
}
|
||||
|
||||
ClauseHeader::ExceptHandler(ExceptHandlerExceptHandler {
|
||||
type_: type_expr,
|
||||
range: _,
|
||||
name: _,
|
||||
body: _,
|
||||
}) => {
|
||||
if let Some(type_expr) = type_expr.as_deref() {
|
||||
visit(type_expr, visitor);
|
||||
}
|
||||
}
|
||||
ClauseHeader::Match(StmtMatch {
|
||||
subject,
|
||||
range: _,
|
||||
cases: _,
|
||||
}) => {
|
||||
visit(subject.as_ref(), visitor);
|
||||
}
|
||||
ClauseHeader::MatchCase(MatchCase {
|
||||
guard,
|
||||
pattern,
|
||||
range: _,
|
||||
body: _,
|
||||
}) => {
|
||||
visit(pattern, visitor);
|
||||
|
||||
if let Some(guard) = guard.as_deref() {
|
||||
visit(guard, visitor);
|
||||
}
|
||||
}
|
||||
ClauseHeader::For(StmtFor {
|
||||
target,
|
||||
iter,
|
||||
range: _,
|
||||
is_async: _,
|
||||
body: _,
|
||||
orelse: _,
|
||||
}) => {
|
||||
visit(target.as_ref(), visitor);
|
||||
visit(iter.as_ref(), visitor);
|
||||
}
|
||||
ClauseHeader::While(StmtWhile {
|
||||
test,
|
||||
range: _,
|
||||
body: _,
|
||||
orelse: _,
|
||||
}) => {
|
||||
visit(test.as_ref(), visitor);
|
||||
}
|
||||
ClauseHeader::With(StmtWith {
|
||||
items,
|
||||
range: _,
|
||||
is_async: _,
|
||||
body: _,
|
||||
}) => {
|
||||
for item in items {
|
||||
visit(item, visitor);
|
||||
}
|
||||
}
|
||||
ClauseHeader::Try(_) | ClauseHeader::TryFinally(_) | ClauseHeader::OrElse(_) => {}
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the range of the first keyword that marks the start of the clause header.
|
||||
fn first_keyword_range(self, source: &str) -> FormatResult<TextRange> {
|
||||
match self {
|
||||
ClauseHeader::Class(header) => {
|
||||
find_keyword(header.start(), SimpleTokenKind::Class, source)
|
||||
}
|
||||
ClauseHeader::Function(header) => {
|
||||
let keyword = if header.is_async {
|
||||
SimpleTokenKind::Async
|
||||
} else {
|
||||
SimpleTokenKind::Def
|
||||
};
|
||||
find_keyword(header.start(), keyword, source)
|
||||
}
|
||||
ClauseHeader::If(header) => find_keyword(header.start(), SimpleTokenKind::If, source),
|
||||
ClauseHeader::ElifElse(ElifElseClause {
|
||||
test: None, range, ..
|
||||
}) => find_keyword(range.start(), SimpleTokenKind::Else, source),
|
||||
ClauseHeader::ElifElse(ElifElseClause {
|
||||
test: Some(_),
|
||||
range,
|
||||
..
|
||||
}) => find_keyword(range.start(), SimpleTokenKind::Elif, source),
|
||||
ClauseHeader::Try(header) => find_keyword(header.start(), SimpleTokenKind::Try, source),
|
||||
ClauseHeader::ExceptHandler(header) => {
|
||||
find_keyword(header.start(), SimpleTokenKind::Except, source)
|
||||
}
|
||||
ClauseHeader::TryFinally(header) => {
|
||||
let last_statement = header
|
||||
.orelse
|
||||
.last()
|
||||
.map(AnyNodeRef::from)
|
||||
.or_else(|| header.handlers.last().map(AnyNodeRef::from))
|
||||
.or_else(|| header.body.last().map(AnyNodeRef::from))
|
||||
.unwrap();
|
||||
|
||||
find_keyword(last_statement.end(), SimpleTokenKind::Finally, source)
|
||||
}
|
||||
ClauseHeader::Match(header) => {
|
||||
find_keyword(header.start(), SimpleTokenKind::Match, source)
|
||||
}
|
||||
ClauseHeader::MatchCase(header) => {
|
||||
find_keyword(header.start(), SimpleTokenKind::Case, source)
|
||||
}
|
||||
ClauseHeader::For(header) => {
|
||||
let keyword = if header.is_async {
|
||||
SimpleTokenKind::Async
|
||||
} else {
|
||||
SimpleTokenKind::For
|
||||
};
|
||||
find_keyword(header.start(), keyword, source)
|
||||
}
|
||||
ClauseHeader::While(header) => {
|
||||
find_keyword(header.start(), SimpleTokenKind::While, source)
|
||||
}
|
||||
ClauseHeader::With(header) => {
|
||||
let keyword = if header.is_async {
|
||||
SimpleTokenKind::Async
|
||||
} else {
|
||||
SimpleTokenKind::With
|
||||
};
|
||||
|
||||
find_keyword(header.start(), keyword, source)
|
||||
}
|
||||
ClauseHeader::OrElse(header) => match header {
|
||||
ElseClause::Try(try_stmt) => {
|
||||
let last_statement = try_stmt
|
||||
.handlers
|
||||
.last()
|
||||
.map(AnyNodeRef::from)
|
||||
.or_else(|| try_stmt.body.last().map(AnyNodeRef::from))
|
||||
.unwrap();
|
||||
|
||||
find_keyword(last_statement.end(), SimpleTokenKind::Else, source)
|
||||
}
|
||||
ElseClause::For(StmtFor { body, .. })
|
||||
| ElseClause::While(StmtWhile { body, .. }) => {
|
||||
find_keyword(body.last().unwrap().end(), SimpleTokenKind::Else, source)
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone)]
|
||||
pub(crate) enum ElseClause<'a> {
|
||||
Try(&'a StmtTry),
|
||||
For(&'a StmtFor),
|
||||
While(&'a StmtWhile),
|
||||
}
|
||||
|
||||
pub(crate) struct FormatClauseHeader<'a, 'ast> {
|
||||
header: ClauseHeader<'a>,
|
||||
/// How to format the clause header
|
||||
formatter: Argument<'a, PyFormatContext<'ast>>,
|
||||
|
||||
/// Leading comments coming before the branch, together with the previous node, if any. Only relevant
|
||||
/// for alternate branches.
|
||||
leading_comments: Option<(&'a [SourceComment], Option<AnyNodeRef<'a>>)>,
|
||||
|
||||
/// The trailing comments coming after the colon.
|
||||
trailing_colon_comment: &'a [SourceComment],
|
||||
}
|
||||
|
||||
/// Formats a clause header, handling the case where the clause header is suppressed and should not be formatted.
|
||||
///
|
||||
/// Calls the `formatter` to format the content of the `header`, except if the `trailing_colon_comment` is a `fmt: skip` suppression comment.
|
||||
/// Takes care of formatting the `trailing_colon_comment` and adds the `:` at the end of the header.
|
||||
pub(crate) fn clause_header<'a, 'ast, Content>(
|
||||
header: ClauseHeader<'a>,
|
||||
trailing_colon_comment: &'a [SourceComment],
|
||||
formatter: &'a Content,
|
||||
) -> FormatClauseHeader<'a, 'ast>
|
||||
where
|
||||
Content: Format<PyFormatContext<'ast>>,
|
||||
{
|
||||
FormatClauseHeader {
|
||||
header,
|
||||
formatter: Argument::new(formatter),
|
||||
leading_comments: None,
|
||||
trailing_colon_comment,
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> FormatClauseHeader<'a, '_> {
|
||||
/// Sets the leading comments that precede an alternate branch.
|
||||
#[must_use]
|
||||
pub(crate) fn with_leading_comments<N>(
|
||||
mut self,
|
||||
comments: &'a [SourceComment],
|
||||
last_node: Option<N>,
|
||||
) -> Self
|
||||
where
|
||||
N: Into<AnyNodeRef<'a>>,
|
||||
{
|
||||
self.leading_comments = Some((comments, last_node.map(Into::into)));
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl<'ast> Format<PyFormatContext<'ast>> for FormatClauseHeader<'_, 'ast> {
|
||||
fn fmt(&self, f: &mut Formatter<PyFormatContext<'ast>>) -> FormatResult<()> {
|
||||
if let Some((leading_comments, last_node)) = self.leading_comments {
|
||||
leading_alternate_branch_comments(leading_comments, last_node).fmt(f)?;
|
||||
}
|
||||
|
||||
if SuppressionKind::has_skip_comment(self.trailing_colon_comment, f.context().source()) {
|
||||
write_suppressed_clause_header(self.header, f)?;
|
||||
} else {
|
||||
f.write_fmt(Arguments::from(&self.formatter))?;
|
||||
text(":").fmt(f)?;
|
||||
}
|
||||
|
||||
trailing_comments(self.trailing_colon_comment).fmt(f)
|
||||
}
|
||||
}
|
||||
|
||||
/// Finds the range of `keyword` starting the search at `start_position`. Expects only comments and `(` between
|
||||
/// the `start_position` and the `keyword` token.
|
||||
fn find_keyword(
|
||||
start_position: TextSize,
|
||||
keyword: SimpleTokenKind,
|
||||
source: &str,
|
||||
) -> FormatResult<TextRange> {
|
||||
let mut tokenizer = SimpleTokenizer::starts_at(start_position, source).skip_trivia();
|
||||
|
||||
match tokenizer.next() {
|
||||
Some(token) if token.kind() == keyword => Ok(token.range()),
|
||||
Some(other) => {
|
||||
debug_assert!(
|
||||
false,
|
||||
"Expected the keyword token {keyword:?} but found the token {other:?} instead."
|
||||
);
|
||||
Err(FormatError::syntax_error(
|
||||
"Expected the keyword token but found another token instead.",
|
||||
))
|
||||
}
|
||||
None => {
|
||||
debug_assert!(
|
||||
false,
|
||||
"Expected the keyword token {keyword:?} but reached the end of the source instead."
|
||||
);
|
||||
Err(FormatError::syntax_error(
|
||||
"Expected the case header keyword token but reached the end of the source instead.",
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the range of the `:` ending the clause header or `Err` if the colon can't be found.
|
||||
fn colon_range(after_keyword_or_condition: TextSize, source: &str) -> FormatResult<TextRange> {
|
||||
let mut tokenizer = SimpleTokenizer::starts_at(after_keyword_or_condition, source)
|
||||
.skip_trivia()
|
||||
.skip_while(|token| token.kind() == SimpleTokenKind::RParen);
|
||||
|
||||
match tokenizer.next() {
|
||||
Some(SimpleToken {
|
||||
kind: SimpleTokenKind::Colon,
|
||||
range,
|
||||
}) => Ok(range),
|
||||
Some(token) => {
|
||||
debug_assert!(false, "Expected the colon marking the end of the case header but found {token:?} instead.");
|
||||
Err(FormatError::syntax_error("Expected colon marking the end of the case header but found another token instead."))
|
||||
}
|
||||
None => {
|
||||
debug_assert!(false, "Expected the colon marking the end of the case header but found the end of the range.");
|
||||
Err(FormatError::syntax_error("Expected the colon marking the end of the case header but found the end of the range."))
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue