From 89f9f044f650c4be2d390133379b0e63c8be2233 Mon Sep 17 00:00:00 2001 From: Micha Reiser Date: Tue, 15 Aug 2023 18:36:18 +0200 Subject: [PATCH] Introduce `clause_header` builder --- .../other/except_handler_except_handler.rs | 79 ++-- .../src/other/match_case.rs | 67 ++-- .../src/statement/clause.rs | 347 ++++++++++++++++++ .../src/statement/mod.rs | 1 + .../src/statement/stmt_class_def.rs | 151 ++++---- .../src/statement/stmt_for.rs | 72 ++-- .../src/statement/stmt_function_def.rs | 201 +++++----- .../src/statement/stmt_if.rs | 83 ++--- .../src/statement/stmt_match.rs | 36 +- .../src/statement/stmt_try.rs | 27 +- .../src/statement/stmt_while.rs | 53 ++- .../src/statement/stmt_with.rs | 97 ++--- crates/ruff_python_formatter/src/verbatim.rs | 311 +--------------- .../format@fmt_skip__or_else.py.snap | 2 +- 14 files changed, 794 insertions(+), 733 deletions(-) create mode 100644 crates/ruff_python_formatter/src/statement/clause.rs diff --git a/crates/ruff_python_formatter/src/other/except_handler_except_handler.rs b/crates/ruff_python_formatter/src/other/except_handler_except_handler.rs index 15ddd1d1fe6587..15e1ef6205e72a 100644 --- a/crates/ruff_python_formatter/src/other/except_handler_except_handler.rs +++ b/crates/ruff_python_formatter/src/other/except_handler_except_handler.rs @@ -1,12 +1,13 @@ -use crate::comments::{trailing_comments, SourceComment, SuppressionKind}; +use ruff_formatter::FormatRuleWithOptions; +use ruff_formatter::{write, Buffer, FormatResult}; +use ruff_python_ast::ExceptHandlerExceptHandler; + +use crate::comments::SourceComment; use crate::expression::maybe_parenthesize_expression; use crate::expression::parentheses::Parenthesize; use crate::prelude::*; -use crate::verbatim::SuppressedClauseHeader; +use crate::statement::clause::{clause_header, ClauseHeader}; use crate::{FormatNodeRule, PyFormatter}; -use ruff_formatter::FormatRuleWithOptions; -use ruff_formatter::{write, Buffer, FormatResult}; -use ruff_python_ast::ExceptHandlerExceptHandler; #[derive(Copy, Clone, Default)] pub enum ExceptHandlerKind { @@ -47,41 +48,45 @@ impl FormatNodeRule for FormatExceptHandlerExceptHan let comments_info = f.context().comments().clone(); let dangling_comments = comments_info.dangling_comments(item); - if SuppressionKind::has_skip_comment(dangling_comments, f.context().source()) { - SuppressedClauseHeader::ExceptHandler(item).fmt(f)?; - } else { - write!( - f, - [ - text("except"), - match self.except_handler_kind { - ExceptHandlerKind::Regular => None, - ExceptHandlerKind::Starred => Some(text("*")), - } - ] - )?; - - if let Some(type_) = type_ { - write!( - f, - [ - space(), - maybe_parenthesize_expression(type_, item, Parenthesize::IfBreaks) - ] - )?; - if let Some(name) = name { - write!(f, [space(), text("as"), space(), name.format()])?; - } - } - - text(":").fmt(f)?; - } - write!( f, [ - trailing_comments(dangling_comments), - block_indent(&body.format()), + clause_header( + ClauseHeader::ExceptHandler(item), + dangling_comments, + &format_with(|f| { + write!( + f, + [ + text("except"), + match self.except_handler_kind { + ExceptHandlerKind::Regular => None, + ExceptHandlerKind::Starred => Some(text("*")), + } + ] + )?; + + if let Some(type_) = type_ { + write!( + f, + [ + space(), + maybe_parenthesize_expression( + type_, + item, + Parenthesize::IfBreaks + ) + ] + )?; + if let Some(name) = name { + write!(f, [space(), text("as"), space(), name.format()])?; + } + } + + Ok(()) + }), + ), + block_indent(&body.format()) ] ) } diff --git a/crates/ruff_python_formatter/src/other/match_case.rs b/crates/ruff_python_formatter/src/other/match_case.rs index a1fe240d91604a..e59bd2c3ae161a 100644 --- a/crates/ruff_python_formatter/src/other/match_case.rs +++ b/crates/ruff_python_formatter/src/other/match_case.rs @@ -1,10 +1,10 @@ use ruff_formatter::{write, Buffer, FormatResult}; use ruff_python_ast::MatchCase; -use crate::comments::{trailing_comments, SourceComment, SuppressionKind}; +use crate::comments::SourceComment; use crate::not_yet_implemented_custom_text; use crate::prelude::*; -use crate::verbatim::SuppressedClauseHeader; +use crate::statement::clause::{clause_header, ClauseHeader}; use crate::{FormatNodeRule, PyFormatter}; #[derive(Default)] @@ -22,40 +22,43 @@ impl FormatNodeRule for FormatMatchCase { let comments = f.context().comments().clone(); let dangling_item_comments = comments.dangling_comments(item); - if SuppressionKind::has_skip_comment(dangling_item_comments, f.context().source()) { - SuppressedClauseHeader::MatchCase(item).fmt(f)?; - } else { - write!( - f, - [ - text("case"), - space(), - format_with(|f: &mut PyFormatter| { - let comments = f.context().comments(); - - for comment in comments.leading_trailing_comments(pattern) { - // This is a lie, but let's go with it. - comment.mark_formatted(); - } + write!( + f, + [ + clause_header( + ClauseHeader::MatchCase(item), + dangling_item_comments, + &format_with(|f| { + write!( + f, + [ + text("case"), + space(), + format_with(|f: &mut PyFormatter| { + let comments = f.context().comments(); - // Replace the whole `format_with` with `pattern.format()` once pattern formatting is implemented. - not_yet_implemented_custom_text("NOT_YET_IMPLEMENTED_Pattern", pattern) - .fmt(f) - }), - ] - )?; + for comment in comments.leading_trailing_comments(pattern) { + // This is a lie, but let's go with it. + comment.mark_formatted(); + } - if let Some(guard) = guard { - write!(f, [space(), text("if"), space(), guard.format()])?; - } + // Replace the whole `format_with` with `pattern.format()` once pattern formatting is implemented. + not_yet_implemented_custom_text( + "NOT_YET_IMPLEMENTED_Pattern", + pattern, + ) + .fmt(f) + }), + ] + )?; - text(":").fmt(f)?; - } + if let Some(guard) = guard { + write!(f, [space(), text("if"), space(), guard.format()])?; + } - write!( - f, - [ - trailing_comments(dangling_item_comments), + Ok(()) + }), + ), block_indent(&body.format()) ] ) diff --git a/crates/ruff_python_formatter/src/statement/clause.rs b/crates/ruff_python_formatter/src/statement/clause.rs new file mode 100644 index 00000000000000..83cb16aaf4eddf --- /dev/null +++ b/crates/ruff_python_formatter/src/statement/clause.rs @@ -0,0 +1,347 @@ +use crate::comments::{ + leading_alternate_branch_comments, trailing_comments, SourceComment, SuppressionKind, +}; +use crate::prelude::*; +use crate::verbatim::write_suppressed_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 { + let keyword_range = self.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(self, visitor: &mut F) + where + F: FnMut(AnyNodeRef), + { + fn visit<'a, N, F>(node: N, visitor: &mut F) + where + N: Into>, + F: FnMut(AnyNodeRef<'a>), + { + visitor(node.into()); + } + + match self { + ClauseHeader::Class(class) => { + if let Some(type_params) = &class.type_params { + visit(type_params.as_ref(), visitor); + } + + if let Some(arguments) = &class.arguments { + visit(arguments.as_ref(), visitor); + } + } + ClauseHeader::Function(function) => { + visit(function.parameters.as_ref(), visitor); + if let Some(type_params) = function.type_params.as_ref() { + visit(type_params, visitor); + } + } + ClauseHeader::If(if_stmt) => { + visit(if_stmt.test.as_ref(), visitor); + } + ClauseHeader::ElifElse(clause) => { + if let Some(test) = clause.test.as_ref() { + visit(test, visitor); + } + } + + ClauseHeader::ExceptHandler(handler) => { + if let Some(ty) = handler.type_.as_deref() { + visit(ty, visitor); + } + } + ClauseHeader::Match(match_stmt) => { + visit(match_stmt.subject.as_ref(), visitor); + } + ClauseHeader::MatchCase(match_case) => { + visit(&match_case.pattern, visitor); + + if let Some(guard) = match_case.guard.as_deref() { + visit(guard, visitor); + } + } + ClauseHeader::For(for_stmt) => { + visit(for_stmt.target.as_ref(), visitor); + visit(for_stmt.iter.as_ref(), visitor); + } + ClauseHeader::While(while_stmt) => { + visit(while_stmt.test.as_ref(), visitor); + } + ClauseHeader::With(with_stmt) => { + for item in &with_stmt.items { + visit(item, visitor); + } + } + ClauseHeader::Try(_) | ClauseHeader::TryFinally(_) | ClauseHeader::OrElse(_) => {} + } + } + + fn keyword_range(self, source: &str) -> FormatResult { + 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) + } + }, + } + } +} + +fn find_keyword( + start_position: TextSize, + keyword: SimpleTokenKind, + source: &str, +) -> FormatResult { + 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.", + )) + } + } +} + +fn colon_range(after_keyword_or_condition: TextSize, source: &str) -> FormatResult { + 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.")) + } + } +} + +#[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>)>, + + /// 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>, +{ + 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( + mut self, + comments: &'a [SourceComment], + last_node: Option, + ) -> Self + where + N: Into>, + { + self.leading_comments = Some((comments, last_node.map(Into::into))); + self + } +} + +impl<'ast> Format> for FormatClauseHeader<'_, 'ast> { + fn fmt(&self, f: &mut Formatter>) -> 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_header(self.header, f)?; + } else { + f.write_fmt(Arguments::from(&self.formatter))?; + text(":").fmt(f)?; + } + + trailing_comments(self.trailing_colon_comment).fmt(f) + } +} diff --git a/crates/ruff_python_formatter/src/statement/mod.rs b/crates/ruff_python_formatter/src/statement/mod.rs index 2eed3ce7443854..7bc0a4c27ff123 100644 --- a/crates/ruff_python_formatter/src/statement/mod.rs +++ b/crates/ruff_python_formatter/src/statement/mod.rs @@ -3,6 +3,7 @@ use ruff_python_ast::Stmt; use crate::prelude::*; +pub(super) mod clause; pub(crate) mod stmt_ann_assign; pub(crate) mod stmt_assert; pub(crate) mod stmt_assign; diff --git a/crates/ruff_python_formatter/src/statement/stmt_class_def.rs b/crates/ruff_python_formatter/src/statement/stmt_class_def.rs index 81ba2c66a4651b..2957950506de36 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_class_def.rs +++ b/crates/ruff_python_formatter/src/statement/stmt_class_def.rs @@ -2,10 +2,11 @@ use ruff_formatter::write; use ruff_python_ast::{Decorator, Ranged, StmtClassDef}; use ruff_python_trivia::lines_after_ignoring_trivia; -use crate::comments::{leading_comments, trailing_comments, SourceComment, SuppressionKind}; +use crate::comments::{leading_comments, trailing_comments, SourceComment}; use crate::prelude::*; use crate::statement::suite::SuiteKind; -use crate::verbatim::SuppressedClauseHeader; + +use crate::statement::clause::{clause_header, ClauseHeader}; use crate::FormatNodeRule; #[derive(Default)] @@ -31,83 +32,81 @@ impl FormatNodeRule for FormatStmtClassDef { let (leading_definition_comments, trailing_definition_comments) = dangling_comments.split_at(trailing_definition_comments_start); - FormatDecorators { - decorators: decorator_list, - leading_definition_comments, - } - .fmt(f)?; - - if SuppressionKind::has_skip_comment(trailing_definition_comments, f.context().source()) { - SuppressedClauseHeader::Class(item).fmt(f)?; - } else { - write!(f, [text("class"), space(), name.format()])?; - - if let Some(type_params) = type_params.as_deref() { - write!(f, [type_params.format()])?; - } - - if let Some(arguments) = arguments.as_deref() { - // Drop empty the arguments node entirely (i.e., remove the parentheses) if it is empty, - // e.g., given: - // ```python - // class A(): - // ... - // ``` - // - // Format as: - // ```python - // class A: - // ... - // ``` - // - // However, preserve any dangling end-of-line comments, e.g., given: - // ```python - // class A( # comment - // ): - // ... - // - // Format as: - // ```python - // class A: # comment - // ... - // ``` - // - // However, the arguments contain any dangling own-line comments, we retain the - // parentheses, e.g., given: - // ```python - // class A( # comment - // # comment - // ): - // ... - // ``` - // - // Format as: - // ```python - // class A( # comment - // # comment - // ): - // ... - // ``` - if arguments.is_empty() - && comments - .dangling_comments(arguments) - .iter() - .all(|comment| comment.line_position().is_end_of_line()) - { - let dangling = comments.dangling_comments(arguments); - write!(f, [trailing_comments(dangling)])?; - } else { - write!(f, [arguments.format()])?; - } - } - - write!(f, [text(":"),])?; - } - write!( f, [ - trailing_comments(trailing_definition_comments), + FormatDecorators { + decorators: decorator_list, + leading_definition_comments, + }, + clause_header( + ClauseHeader::Class(item), + trailing_definition_comments, + &format_with(|f| { + write!(f, [text("class"), space(), name.format()])?; + + if let Some(type_params) = type_params.as_deref() { + write!(f, [type_params.format()])?; + } + + if let Some(arguments) = arguments.as_deref() { + // Drop empty the arguments node entirely (i.e., remove the parentheses) if it is empty, + // e.g., given: + // ```python + // class A(): + // ... + // ``` + // + // Format as: + // ```python + // class A: + // ... + // ``` + // + // However, preserve any dangling end-of-line comments, e.g., given: + // ```python + // class A( # comment + // ): + // ... + // + // Format as: + // ```python + // class A: # comment + // ... + // ``` + // + // However, the arguments contain any dangling own-line comments, we retain the + // parentheses, e.g., given: + // ```python + // class A( # comment + // # comment + // ): + // ... + // ``` + // + // Format as: + // ```python + // class A( # comment + // # comment + // ): + // ... + // ``` + if arguments.is_empty() + && comments + .dangling_comments(arguments) + .iter() + .all(|comment| comment.line_position().is_end_of_line()) + { + let dangling = comments.dangling_comments(arguments); + write!(f, [trailing_comments(dangling)])?; + } else { + write!(f, [arguments.format()])?; + } + } + + Ok(()) + }), + ), block_indent(&body.format().with_options(SuiteKind::Class)) ] ) diff --git a/crates/ruff_python_formatter/src/statement/stmt_for.rs b/crates/ruff_python_formatter/src/statement/stmt_for.rs index 19cc162aaa1390..e6dea2b0e770df 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_for.rs +++ b/crates/ruff_python_formatter/src/statement/stmt_for.rs @@ -1,14 +1,12 @@ use ruff_formatter::{format_args, write}; use ruff_python_ast::{Expr, Ranged, Stmt, StmtFor}; -use crate::comments::{ - leading_alternate_branch_comments, trailing_comments, SourceComment, SuppressionKind, -}; +use crate::comments::SourceComment; use crate::expression::expr_tuple::TupleParentheses; use crate::expression::maybe_parenthesize_expression; use crate::expression::parentheses::Parenthesize; use crate::prelude::*; -use crate::verbatim::{OrElseParent, SuppressedClauseHeader}; +use crate::statement::clause::{clause_header, ClauseHeader, ElseClause}; use crate::FormatNodeRule; #[derive(Debug)] @@ -49,32 +47,28 @@ impl FormatNodeRule for FormatStmtFor { let (trailing_condition_comments, or_else_comments) = dangling_comments.split_at(or_else_comments_start); - if SuppressionKind::has_skip_comment(trailing_condition_comments, f.context().source()) { - SuppressedClauseHeader::For(item).fmt(f)?; - } else { - write!( - f, - [ - is_async.then_some(format_args![text("async"), space()]), - text("for"), - space(), - ExprTupleWithoutParentheses(target), - space(), - text("in"), - space(), - maybe_parenthesize_expression(iter, item, Parenthesize::IfBreaks), - text(":"), - ] - )?; - } + clause_header( + ClauseHeader::For(item), + trailing_condition_comments, + &format_with(|f| { + write!( + f, + [ + is_async.then_some(format_args![text("async"), space()]), + text("for"), + space(), + ExprTupleWithoutParentheses(target), + space(), + text("in"), + space(), + maybe_parenthesize_expression(iter, item, Parenthesize::IfBreaks), + ] + ) + }), + ) + .fmt(f)?; - write!( - f, - [ - trailing_comments(trailing_condition_comments), - block_indent(&body.format()) - ] - )?; + block_indent(&body.format()).fmt(f)?; if orelse.is_empty() { debug_assert!(or_else_comments.is_empty()); @@ -85,19 +79,15 @@ impl FormatNodeRule for FormatStmtFor { or_else_comments.partition_point(|comment| comment.line_position().is_own_line()); let (leading, trailing) = or_else_comments.split_at(trailing_start); - leading_alternate_branch_comments(leading, body.last()).fmt(f)?; - - if SuppressionKind::has_skip_comment(trailing_condition_comments, f.context().source()) - { - SuppressedClauseHeader::OrElse(OrElseParent::For(item)).fmt(f)?; - } else { - text("else:").fmt(f)?; - } + clause_header( + ClauseHeader::OrElse(ElseClause::For(item)), + trailing, + &text("else"), + ) + .with_leading_comments(leading, body.last()) + .fmt(f)?; - write!( - f, - [trailing_comments(trailing), block_indent(&orelse.format())] - )?; + block_indent(&orelse.format()).fmt(f)?; } Ok(()) diff --git a/crates/ruff_python_formatter/src/statement/stmt_function_def.rs b/crates/ruff_python_formatter/src/statement/stmt_function_def.rs index 6afc97597a7c52..394d1721a644a8 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_function_def.rs +++ b/crates/ruff_python_formatter/src/statement/stmt_function_def.rs @@ -2,13 +2,13 @@ use ruff_formatter::write; use ruff_python_ast::{Parameters, Ranged, StmtFunctionDef}; use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer}; -use crate::comments::{trailing_comments, SourceComment, SuppressionKind}; +use crate::comments::SourceComment; use crate::expression::maybe_parenthesize_expression; use crate::expression::parentheses::{Parentheses, Parenthesize}; use crate::prelude::*; +use crate::statement::clause::{clause_header, ClauseHeader}; use crate::statement::stmt_class_def::FormatDecorators; use crate::statement::suite::SuiteKind; -use crate::verbatim::SuppressedClauseHeader; use crate::FormatNodeRule; #[derive(Default)] @@ -36,106 +36,109 @@ impl FormatNodeRule for FormatStmtFunctionDef { let (leading_definition_comments, trailing_definition_comments) = dangling_comments.split_at(trailing_definition_comments_start); - FormatDecorators { - decorators: decorator_list, - leading_definition_comments, - } - .fmt(f)?; - - if SuppressionKind::has_skip_comment(trailing_definition_comments, f.context().source()) { - SuppressedClauseHeader::Function(item).fmt(f)?; - } else { - if *is_async { - write!(f, [text("async"), space()])?; - } - - write!(f, [text("def"), space(), name.format()])?; - - if let Some(type_params) = type_params.as_ref() { - write!(f, [type_params.format()])?; - } - - let format_inner = format_with(|f: &mut PyFormatter| { - write!(f, [parameters.format()])?; - - if let Some(return_annotation) = returns.as_ref() { - write!(f, [space(), text("->"), space()])?; - - if return_annotation.is_tuple_expr() { - write!( - f, - [return_annotation.format().with_options(Parentheses::Never)] - )?; - } else if comments.has_trailing_comments(return_annotation.as_ref()) { - // Intentionally parenthesize any return annotations with trailing comments. - // This avoids an instability in cases like: - // ```python - // def double( - // a: int - // ) -> ( - // int # Hello - // ): - // pass - // ``` - // If we allow this to break, it will be formatted as follows: - // ```python - // def double( - // a: int - // ) -> int: # Hello - // pass - // ``` - // On subsequent formats, the `# Hello` will be interpreted as a dangling - // comment on a function, yielding: - // ```python - // def double(a: int) -> int: # Hello - // pass - // ``` - // Ideally, we'd reach that final formatting in a single pass, but doing so - // requires that the parent be aware of how the child is formatted, which - // is challenging. As a compromise, we break those expressions to avoid an - // instability. - write!( - f, - [return_annotation.format().with_options(Parentheses::Always)] - )?; - } else { - write!( - f, - [maybe_parenthesize_expression( - return_annotation, - item, - if empty_parameters(parameters, f.context().source()) { - // If the parameters are empty, add parentheses if the return annotation - // breaks at all. - Parenthesize::IfBreaksOrIfRequired - } else { - // Otherwise, use our normal rules for parentheses, which allows us to break - // like: - // ```python - // def f( - // x, - // ) -> Tuple[ - // int, - // int, - // ]: - // ... - // ``` - Parenthesize::IfBreaks - }, - )] - )?; - } - } - Ok(()) - }); - - write!(f, [group(&format_inner), text(":"),])?; - } - write!( f, [ - trailing_comments(trailing_definition_comments), + FormatDecorators { + decorators: decorator_list, + leading_definition_comments, + }, + clause_header( + ClauseHeader::Function(item), + trailing_definition_comments, + &format_with(|f| { + if *is_async { + write!(f, [text("async"), space()])?; + } + + write!(f, [text("def"), space(), name.format()])?; + + if let Some(type_params) = type_params.as_ref() { + write!(f, [type_params.format()])?; + } + + let format_inner = format_with(|f: &mut PyFormatter| { + write!(f, [parameters.format()])?; + + if let Some(return_annotation) = returns.as_ref() { + write!(f, [space(), text("->"), space()])?; + + if return_annotation.is_tuple_expr() { + write!( + f, + [return_annotation + .format() + .with_options(Parentheses::Never)] + )?; + } else if comments.has_trailing_comments(return_annotation.as_ref()) + { + // Intentionally parenthesize any return annotations with trailing comments. + // This avoids an instability in cases like: + // ```python + // def double( + // a: int + // ) -> ( + // int # Hello + // ): + // pass + // ``` + // If we allow this to break, it will be formatted as follows: + // ```python + // def double( + // a: int + // ) -> int: # Hello + // pass + // ``` + // On subsequent formats, the `# Hello` will be interpreted as a dangling + // comment on a function, yielding: + // ```python + // def double(a: int) -> int: # Hello + // pass + // ``` + // Ideally, we'd reach that final formatting in a single pass, but doing so + // requires that the parent be aware of how the child is formatted, which + // is challenging. As a compromise, we break those expressions to avoid an + // instability. + write!( + f, + [return_annotation + .format() + .with_options(Parentheses::Always)] + )?; + } else { + write!( + f, + [maybe_parenthesize_expression( + return_annotation, + item, + if empty_parameters(parameters, f.context().source()) { + // If the parameters are empty, add parentheses if the return annotation + // breaks at all. + Parenthesize::IfBreaksOrIfRequired + } else { + // Otherwise, use our normal rules for parentheses, which allows us to break + // like: + // ```python + // def f( + // x, + // ) -> Tuple[ + // int, + // int, + // ]: + // ... + // ``` + Parenthesize::IfBreaks + }, + )] + )?; + } + } + Ok(()) + }); + + group(&format_inner).fmt(f) + }), + ), block_indent(&body.format().with_options(SuiteKind::Function)) ] ) diff --git a/crates/ruff_python_formatter/src/statement/stmt_if.rs b/crates/ruff_python_formatter/src/statement/stmt_if.rs index 895602a3a8cf36..ed4d15b9c08210 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_if.rs +++ b/crates/ruff_python_formatter/src/statement/stmt_if.rs @@ -1,14 +1,13 @@ -use crate::comments::{ - leading_alternate_branch_comments, trailing_comments, SourceComment, SuppressionKind, -}; +use ruff_formatter::write; +use ruff_python_ast::node::AnyNodeRef; +use ruff_python_ast::{ElifElseClause, StmtIf}; + +use crate::comments::SourceComment; use crate::expression::maybe_parenthesize_expression; use crate::expression::parentheses::Parenthesize; use crate::prelude::*; -use crate::verbatim::SuppressedClauseHeader; +use crate::statement::clause::{clause_header, ClauseHeader}; use crate::FormatNodeRule; -use ruff_formatter::write; -use ruff_python_ast::node::AnyNodeRef; -use ruff_python_ast::{ElifElseClause, StmtIf}; #[derive(Default)] pub struct FormatStmtIf; @@ -25,24 +24,23 @@ impl FormatNodeRule for FormatStmtIf { let comments = f.context().comments().clone(); let trailing_colon_comment = comments.dangling_comments(item); - if SuppressionKind::has_skip_comment(trailing_colon_comment, f.context().source()) { - SuppressedClauseHeader::If(item).fmt(f)?; - } else { - write!( - f, - [ - text("if"), - space(), - maybe_parenthesize_expression(test, item, Parenthesize::IfBreaks), - text(":"), - ] - )?; - } - write!( f, [ - trailing_comments(trailing_colon_comment), + clause_header( + ClauseHeader::If(item), + trailing_colon_comment, + &format_with(|f| { + write!( + f, + [ + text("if"), + space(), + maybe_parenthesize_expression(test, item, Parenthesize::IfBreaks), + ] + ) + }), + ), block_indent(&body.format()) ] )?; @@ -83,31 +81,28 @@ pub(crate) fn format_elif_else_clause( let trailing_colon_comment = comments.dangling_comments(item); let leading_comments = comments.leading_comments(item); - leading_alternate_branch_comments(leading_comments, last_node).fmt(f)?; - - if SuppressionKind::has_skip_comment(trailing_colon_comment, f.context().source()) { - SuppressedClauseHeader::ElifElse(item).fmt(f)?; - } else { - if let Some(test) = test { - write!( - f, - [ - text("elif"), - space(), - maybe_parenthesize_expression(test, item, Parenthesize::IfBreaks), - ] - )?; - } else { - text("else").fmt(f)?; - } - - text(":").fmt(f)?; - } - write!( f, [ - trailing_comments(trailing_colon_comment), + clause_header( + ClauseHeader::ElifElse(item), + trailing_colon_comment, + &format_with(|f| { + if let Some(test) = test { + write!( + f, + [ + text("elif"), + space(), + maybe_parenthesize_expression(test, item, Parenthesize::IfBreaks), + ] + ) + } else { + text("else").fmt(f) + } + }), + ) + .with_leading_comments(leading_comments, last_node), block_indent(&body.format()) ] ) diff --git a/crates/ruff_python_formatter/src/statement/stmt_match.rs b/crates/ruff_python_formatter/src/statement/stmt_match.rs index 27ba2b0e832ad7..a684718d4ab63f 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_match.rs +++ b/crates/ruff_python_formatter/src/statement/stmt_match.rs @@ -1,14 +1,12 @@ use ruff_formatter::{format_args, write}; use ruff_python_ast::StmtMatch; -use crate::comments::{ - leading_alternate_branch_comments, trailing_comments, SourceComment, SuppressionKind, -}; +use crate::comments::{leading_alternate_branch_comments, SourceComment}; use crate::context::{NodeLevel, WithNodeLevel}; use crate::expression::maybe_parenthesize_expression; use crate::expression::parentheses::Parenthesize; use crate::prelude::*; -use crate::verbatim::SuppressedClauseHeader; +use crate::statement::clause::{clause_header, ClauseHeader}; use crate::FormatNodeRule; #[derive(Default)] @@ -28,21 +26,21 @@ impl FormatNodeRule for FormatStmtMatch { // There can be at most one dangling comment after the colon in a match statement. debug_assert!(dangling_item_comments.len() <= 1); - if SuppressionKind::has_skip_comment(dangling_item_comments, f.context().source()) { - SuppressedClauseHeader::Match(item).fmt(f)?; - } else { - write!( - f, - [ - text("match"), - space(), - maybe_parenthesize_expression(subject, item, Parenthesize::IfBreaks), - text(":"), - ] - )?; - } - - trailing_comments(dangling_item_comments).fmt(f)?; + clause_header( + ClauseHeader::Match(item), + dangling_item_comments, + &format_with(|f| { + write!( + f, + [ + text("match"), + space(), + maybe_parenthesize_expression(subject, item, Parenthesize::IfBreaks), + ] + ) + }), + ) + .fmt(f)?; let mut cases_iter = cases.iter(); let Some(first) = cases_iter.next() else { diff --git a/crates/ruff_python_formatter/src/statement/stmt_try.rs b/crates/ruff_python_formatter/src/statement/stmt_try.rs index f80f465652bd0f..17263898a42b63 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_try.rs +++ b/crates/ruff_python_formatter/src/statement/stmt_try.rs @@ -2,12 +2,12 @@ use ruff_formatter::{write, FormatRuleWithOptions}; use ruff_python_ast::{ExceptHandler, Ranged, StmtTry}; use crate::comments; -use crate::comments::{leading_alternate_branch_comments, trailing_comments}; -use crate::comments::{SourceComment, SuppressionKind}; +use crate::comments::leading_alternate_branch_comments; +use crate::comments::SourceComment; use crate::other::except_handler_except_handler::ExceptHandlerKind; use crate::prelude::*; +use crate::statement::clause::{clause_header, ClauseHeader, ElseClause}; use crate::statement::{FormatRefWithRule, Stmt}; -use crate::verbatim::{OrElseParent, SuppressedClauseHeader}; use crate::{FormatNodeRule, PyFormatter}; #[derive(Default)] @@ -127,23 +127,17 @@ fn format_case<'a>( let (leading_case_comments, trailing_case_comments) = case_comments.split_at(partition_point); - leading_alternate_branch_comments(leading_case_comments, previous_node).fmt(f)?; + let header = match kind { + CaseKind::Try => ClauseHeader::Try(try_statement), + CaseKind::Else => ClauseHeader::OrElse(ElseClause::Try(try_statement)), + CaseKind::Finally => ClauseHeader::TryFinally(try_statement), + }; - if SuppressionKind::has_skip_comment(trailing_case_comments, f.context().source()) { - let header = match kind { - CaseKind::Try => SuppressedClauseHeader::Try(try_statement), - CaseKind::Else => SuppressedClauseHeader::OrElse(OrElseParent::Try(try_statement)), - CaseKind::Finally => SuppressedClauseHeader::TryFinally(try_statement), - }; - - header.fmt(f)?; - } else { - write!(f, [text(kind.keyword()), text(":")])?; - } write!( f, [ - trailing_comments(trailing_case_comments), + clause_header(header, trailing_case_comments, &text(kind.keyword())) + .with_leading_comments(leading_case_comments, previous_node), block_indent(&body.format()) ] )?; @@ -153,6 +147,7 @@ fn format_case<'a>( }) } +#[derive(Copy, Clone)] enum CaseKind { Try, Else, diff --git a/crates/ruff_python_formatter/src/statement/stmt_while.rs b/crates/ruff_python_formatter/src/statement/stmt_while.rs index a99e83c70d2073..67266356b032be 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_while.rs +++ b/crates/ruff_python_formatter/src/statement/stmt_while.rs @@ -2,13 +2,11 @@ use ruff_formatter::write; use ruff_python_ast::node::AstNode; use ruff_python_ast::{Ranged, Stmt, StmtWhile}; -use crate::comments::{ - leading_alternate_branch_comments, trailing_comments, SourceComment, SuppressionKind, -}; +use crate::comments::SourceComment; use crate::expression::maybe_parenthesize_expression; use crate::expression::parentheses::Parenthesize; use crate::prelude::*; -use crate::verbatim::{OrElseParent, SuppressedClauseHeader}; +use crate::statement::clause::{clause_header, ClauseHeader, ElseClause}; use crate::FormatNodeRule; #[derive(Default)] @@ -33,24 +31,23 @@ impl FormatNodeRule for FormatStmtWhile { let (trailing_condition_comments, or_else_comments) = dangling_comments.split_at(or_else_comments_start); - if SuppressionKind::has_skip_comment(trailing_condition_comments, f.context().source()) { - SuppressedClauseHeader::While(item).fmt(f)?; - } else { - write!( - f, - [ - text("while"), - space(), - maybe_parenthesize_expression(test, item, Parenthesize::IfBreaks), - text(":"), - ] - )?; - } - write!( f, [ - trailing_comments(trailing_condition_comments), + clause_header( + ClauseHeader::While(item), + trailing_condition_comments, + &format_with(|f| { + write!( + f, + [ + text("while"), + space(), + maybe_parenthesize_expression(test, item, Parenthesize::IfBreaks), + ] + ) + }), + ), block_indent(&body.format()) ] )?; @@ -62,17 +59,17 @@ impl FormatNodeRule for FormatStmtWhile { or_else_comments.partition_point(|comment| comment.line_position().is_own_line()); let (leading, trailing) = or_else_comments.split_at(trailing_start); - leading_alternate_branch_comments(leading, body.last()).fmt(f)?; - - if SuppressionKind::has_skip_comment(trailing, f.context().source()) { - SuppressedClauseHeader::OrElse(OrElseParent::While(item)).fmt(f)?; - } else { - text("else:").fmt(f)?; - } - write!( f, - [trailing_comments(trailing), block_indent(&orelse.format())] + [ + clause_header( + ClauseHeader::OrElse(ElseClause::While(item)), + trailing, + &text("else") + ) + .with_leading_comments(leading, body.last()), + block_indent(&orelse.format()) + ] )?; } diff --git a/crates/ruff_python_formatter/src/statement/stmt_with.rs b/crates/ruff_python_formatter/src/statement/stmt_with.rs index 18003407ed7ced..b96e4f75352d11 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_with.rs +++ b/crates/ruff_python_formatter/src/statement/stmt_with.rs @@ -4,12 +4,12 @@ use ruff_python_ast::{Ranged, StmtWith}; use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer}; use ruff_text_size::TextRange; -use crate::comments::{trailing_comments, SourceComment, SuppressionKind}; +use crate::comments::SourceComment; use crate::expression::parentheses::{ in_parentheses_only_soft_line_break_or_space, optional_parentheses, parenthesized, }; use crate::prelude::*; -use crate::verbatim::SuppressedClauseHeader; +use crate::statement::clause::{clause_header, ClauseHeader}; use crate::FormatNodeRule; #[derive(Default)] @@ -39,56 +39,57 @@ impl FormatNodeRule for FormatStmtWith { }); let (parenthesized_comments, colon_comments) = dangling_comments.split_at(partition_point); - if SuppressionKind::has_skip_comment(colon_comments, f.context().source()) { - SuppressedClauseHeader::With(item).fmt(f)?; - } else { - write!( - f, - [ - item.is_async - .then_some(format_args![text("async"), space()]), - text("with"), - space() - ] - )?; - - if !parenthesized_comments.is_empty() { - let joined = format_with(|f: &mut PyFormatter| { - f.join_comma_separated(item.body.first().unwrap().start()) - .nodes(&item.items) - .finish() - }); + write!( + f, + [ + clause_header( + ClauseHeader::With(item), + colon_comments, + &format_with(|f| { + write!( + f, + [ + item.is_async + .then_some(format_args![text("async"), space()]), + text("with"), + space() + ] + )?; - parenthesized("(", &joined, ")") - .with_dangling_comments(parenthesized_comments) - .fmt(f)?; - } else if are_with_items_parenthesized(item, f.context())? { - optional_parentheses(&format_with(|f| { - let mut joiner = f.join_comma_separated(item.body.first().unwrap().start()); + if !parenthesized_comments.is_empty() { + let joined = format_with(|f: &mut PyFormatter| { + f.join_comma_separated(item.body.first().unwrap().start()) + .nodes(&item.items) + .finish() + }); - for item in &item.items { - joiner.entry_with_line_separator( - item, - &item.format(), - in_parentheses_only_soft_line_break_or_space(), - ); - } - joiner.finish() - })) - .fmt(f)?; - } else { - f.join_with(format_args![text(","), space()]) - .entries(item.items.iter().formatted()) - .finish()?; - } + parenthesized("(", &joined, ")") + .with_dangling_comments(parenthesized_comments) + .fmt(f)?; + } else if are_with_items_parenthesized(item, f.context())? { + optional_parentheses(&format_with(|f| { + let mut joiner = + f.join_comma_separated(item.body.first().unwrap().start()); - text(":").fmt(f)?; - } + for item in &item.items { + joiner.entry_with_line_separator( + item, + &item.format(), + in_parentheses_only_soft_line_break_or_space(), + ); + } + joiner.finish() + })) + .fmt(f)?; + } else { + f.join_with(format_args![text(","), space()]) + .entries(item.items.iter().formatted()) + .finish()?; + } - write!( - f, - [ - trailing_comments(colon_comments), + Ok(()) + }) + ), block_indent(&item.body.format()) ] ) diff --git a/crates/ruff_python_formatter/src/verbatim.rs b/crates/ruff_python_formatter/src/verbatim.rs index 533cc345d230ef..c551c263248f3c 100644 --- a/crates/ruff_python_formatter/src/verbatim.rs +++ b/crates/ruff_python_formatter/src/verbatim.rs @@ -5,19 +5,17 @@ use unicode_width::UnicodeWidthStr; use ruff_formatter::{write, FormatError}; use ruff_python_ast::node::AnyNodeRef; -use ruff_python_ast::{ - ElifElseClause, ExceptHandlerExceptHandler, MatchCase, Ranged, Stmt, StmtClassDef, StmtFor, - StmtFunctionDef, StmtIf, StmtMatch, StmtTry, StmtWhile, StmtWith, -}; +use ruff_python_ast::{Ranged, Stmt}; use ruff_python_parser::lexer::{lex_starts_at, LexResult}; use ruff_python_parser::{Mode, Tok}; -use ruff_python_trivia::{lines_before, SimpleToken, SimpleTokenKind, SimpleTokenizer}; +use ruff_python_trivia::lines_before; use ruff_source_file::Locator; use ruff_text_size::{TextRange, TextSize}; use crate::comments::format::{empty_lines, format_comment}; use crate::comments::{leading_comments, trailing_comments, SourceComment}; use crate::prelude::*; +use crate::statement::clause::ClauseHeader; use crate::statement::suite::SuiteChildStatement; /// Disables formatting for all statements between the `first_suppressed` that has a leading `fmt: off` comment @@ -933,293 +931,22 @@ impl Format> for FormatSuppressedNode<'_> { } } -#[derive(Copy, Clone)] -pub(crate) enum SuppressedClauseHeader<'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(OrElseParent<'a>), -} - -impl<'a> SuppressedClauseHeader<'a> { - fn range(self, source: &str) -> FormatResult { - let keyword_range = self.keyword_range(source)?; - - let mut last_child_end = None; - - self.visit_children(&mut |child| last_child_end = Some(child.end())); - - let end = match self { - SuppressedClauseHeader::Class(class) => { - Some(last_child_end.unwrap_or(class.name.end())) - } - SuppressedClauseHeader::Function(function) => { - Some(last_child_end.unwrap_or(function.name.end())) - } - SuppressedClauseHeader::ElifElse(_) - | SuppressedClauseHeader::Try(_) - | SuppressedClauseHeader::If(_) - | SuppressedClauseHeader::TryFinally(_) - | SuppressedClauseHeader::Match(_) - | SuppressedClauseHeader::MatchCase(_) - | SuppressedClauseHeader::For(_) - | SuppressedClauseHeader::While(_) - | SuppressedClauseHeader::With(_) - | SuppressedClauseHeader::OrElse(_) => last_child_end, - - SuppressedClauseHeader::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())) - } - - // Needs to know child nodes. - // Could use normal lexer to get header end since we know that it is rare - fn visit_children(self, visitor: &mut F) - where - F: FnMut(AnyNodeRef), - { - fn visit<'a, N, F>(node: N, visitor: &mut F) - where - N: Into>, - F: FnMut(AnyNodeRef<'a>), - { - visitor(node.into()); - } - - match self { - SuppressedClauseHeader::Class(class) => { - if let Some(type_params) = &class.type_params { - visit(type_params.as_ref(), visitor); - } - - if let Some(arguments) = &class.arguments { - visit(arguments.as_ref(), visitor); - } - } - SuppressedClauseHeader::Function(function) => { - visit(function.parameters.as_ref(), visitor); - if let Some(type_params) = function.type_params.as_ref() { - visit(type_params, visitor); - } - } - SuppressedClauseHeader::If(if_stmt) => { - visit(if_stmt.test.as_ref(), visitor); - } - SuppressedClauseHeader::ElifElse(clause) => { - if let Some(test) = clause.test.as_ref() { - visit(test, visitor); - } - } - - SuppressedClauseHeader::ExceptHandler(handler) => { - if let Some(ty) = handler.type_.as_deref() { - visit(ty, visitor); - } - } - SuppressedClauseHeader::Match(match_stmt) => { - visit(match_stmt.subject.as_ref(), visitor); - } - SuppressedClauseHeader::MatchCase(match_case) => { - visit(&match_case.pattern, visitor); - - if let Some(guard) = match_case.guard.as_deref() { - visit(guard, visitor); - } - } - SuppressedClauseHeader::For(for_stmt) => { - visit(for_stmt.target.as_ref(), visitor); - visit(for_stmt.iter.as_ref(), visitor); - } - SuppressedClauseHeader::While(while_stmt) => { - visit(while_stmt.test.as_ref(), visitor); - } - SuppressedClauseHeader::With(with_stmt) => { - for item in &with_stmt.items { - visit(item, visitor); - } - } - SuppressedClauseHeader::Try(_) - | SuppressedClauseHeader::TryFinally(_) - | SuppressedClauseHeader::OrElse(_) => {} - } - } - - fn keyword_range(self, source: &str) -> FormatResult { - match self { - SuppressedClauseHeader::Class(header) => { - find_keyword(header.start(), SimpleTokenKind::Class, source) - } - SuppressedClauseHeader::Function(header) => { - let keyword = if header.is_async { - SimpleTokenKind::Async - } else { - SimpleTokenKind::Def - }; - find_keyword(header.start(), keyword, source) - } - SuppressedClauseHeader::If(header) => { - find_keyword(header.start(), SimpleTokenKind::If, source) - } - SuppressedClauseHeader::ElifElse(ElifElseClause { - test: None, range, .. - }) => find_keyword(range.start(), SimpleTokenKind::Else, source), - SuppressedClauseHeader::ElifElse(ElifElseClause { - test: Some(_), - range, - .. - }) => find_keyword(range.start(), SimpleTokenKind::Elif, source), - SuppressedClauseHeader::Try(header) => { - find_keyword(header.start(), SimpleTokenKind::Try, source) - } - SuppressedClauseHeader::ExceptHandler(header) => { - find_keyword(header.start(), SimpleTokenKind::Except, source) - } - SuppressedClauseHeader::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) - } - SuppressedClauseHeader::Match(header) => { - find_keyword(header.start(), SimpleTokenKind::Match, source) - } - SuppressedClauseHeader::MatchCase(header) => { - find_keyword(header.start(), SimpleTokenKind::Case, source) - } - SuppressedClauseHeader::For(header) => { - let keyword = if header.is_async { - SimpleTokenKind::Async - } else { - SimpleTokenKind::For - }; - find_keyword(header.start(), keyword, source) - } - SuppressedClauseHeader::While(header) => { - find_keyword(header.start(), SimpleTokenKind::While, source) - } - SuppressedClauseHeader::With(header) => { - let keyword = if header.is_async { - SimpleTokenKind::Async - } else { - SimpleTokenKind::With - }; - - find_keyword(header.start(), keyword, source) - } - SuppressedClauseHeader::OrElse(header) => match header { - OrElseParent::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) - } - OrElseParent::For(StmtFor { body, .. }) - | OrElseParent::While(StmtWhile { body, .. }) => { - find_keyword(body.last().unwrap().end(), SimpleTokenKind::Else, source) - } - }, - } - } -} - -fn find_keyword( - start_position: TextSize, - keyword: SimpleTokenKind, - source: &str, -) -> FormatResult { - 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.", - )) - } - } -} - -fn colon_range(after_keyword_or_condition: TextSize, source: &str) -> FormatResult { - 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.")) - } - } -} - -#[derive(Copy, Clone)] -pub(crate) enum OrElseParent<'a> { - Try(&'a StmtTry), - For(&'a StmtFor), - While(&'a StmtWhile), -} - -impl Format> for SuppressedClauseHeader<'_> { - #[cold] - fn fmt(&self, f: &mut Formatter>) -> FormatResult<()> { - // Write the outer comments and format the node as verbatim - write!( - f, - [verbatim_text( - self.range(f.context().source())?, - ContainsNewlines::Detect - ),] - )?; +#[cold] +pub(crate) fn write_suppressed_header( + header: ClauseHeader, + f: &mut PyFormatter, +) -> FormatResult<()> { + // Write the outer comments and format the node as verbatim + write!( + f, + [verbatim_text( + header.range(f.context().source())?, + ContainsNewlines::Detect + ),] + )?; - let comments = f.context().comments(); - self.visit_children(&mut |child| comments.mark_verbatim_node_comments_formatted(child)); + let comments = f.context().comments(); + header.visit(&mut |child| comments.mark_verbatim_node_comments_formatted(child)); - Ok(()) - } + Ok(()) } diff --git a/crates/ruff_python_formatter/tests/snapshots/format@fmt_skip__or_else.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@fmt_skip__or_else.py.snap index 21bad10de4353c..4293169e6bb0bd 100644 --- a/crates/ruff_python_formatter/tests/snapshots/format@fmt_skip__or_else.py.snap +++ b/crates/ruff_python_formatter/tests/snapshots/format@fmt_skip__or_else.py.snap @@ -45,7 +45,7 @@ for item in container: process(item) break # leading comment -else: # fmt: skip +else : # fmt: skip # Didn't find anything.. not_found_in_container()