diff --git a/crates/ruff_python_formatter/src/format/excepthandler.rs b/crates/ruff_python_formatter/src/format/excepthandler.rs new file mode 100644 index 0000000000..8c59d8ee2b --- /dev/null +++ b/crates/ruff_python_formatter/src/format/excepthandler.rs @@ -0,0 +1,46 @@ +use ruff_formatter::prelude::*; +use ruff_formatter::write; +use ruff_text_size::TextSize; + +use crate::context::ASTFormatContext; +use crate::cst::{Excepthandler, ExcepthandlerKind}; +use crate::format::builders::block; +use crate::shared_traits::AsFormat; + +pub struct FormatExcepthandler<'a> { + item: &'a Excepthandler, +} + +impl AsFormat> for Excepthandler { + type Format<'a> = FormatExcepthandler<'a>; + + fn format(&self) -> Self::Format<'_> { + FormatExcepthandler { item: self } + } +} + +impl Format> for FormatExcepthandler<'_> { + fn fmt(&self, f: &mut Formatter) -> FormatResult<()> { + let ExcepthandlerKind::ExceptHandler { type_, name, body } = &self.item.node; + + write!(f, [text("except")])?; + if let Some(type_) = &type_ { + write!(f, [space(), type_.format()])?; + if let Some(name) = &name { + write!( + f, + [ + space(), + text("as"), + space(), + dynamic_text(name, TextSize::default()), + ] + )?; + } + } + write!(f, [text(":")])?; + write!(f, [block_indent(&block(body))])?; + + Ok(()) + } +} diff --git a/crates/ruff_python_formatter/src/format/mod.rs b/crates/ruff_python_formatter/src/format/mod.rs index 9437853cc1..ca673053b3 100644 --- a/crates/ruff_python_formatter/src/format/mod.rs +++ b/crates/ruff_python_formatter/src/format/mod.rs @@ -5,6 +5,7 @@ mod boolop; pub mod builders; mod cmpop; mod comprehension; +mod excepthandler; mod expr; mod helpers; mod operator; diff --git a/crates/ruff_python_formatter/src/format/stmt.rs b/crates/ruff_python_formatter/src/format/stmt.rs index e088d1624b..a8742e72b5 100644 --- a/crates/ruff_python_formatter/src/format/stmt.rs +++ b/crates/ruff_python_formatter/src/format/stmt.rs @@ -6,7 +6,9 @@ use ruff_text_size::TextSize; use crate::builders::literal; use crate::context::ASTFormatContext; -use crate::cst::{Alias, Arguments, Expr, ExprKind, Keyword, Stmt, StmtKind, Withitem}; +use crate::cst::{ + Alias, Arguments, Excepthandler, Expr, ExprKind, Keyword, Stmt, StmtKind, Withitem, +}; use crate::format::builders::{block, join_names}; use crate::format::helpers::is_self_closing; use crate::shared_traits::AsFormat; @@ -419,12 +421,54 @@ fn format_raise( fn format_return( f: &mut Formatter>, + stmt: &Stmt, value: Option<&Expr>, ) -> FormatResult<()> { write!(f, [text("return")])?; if let Some(value) = value { write!(f, [space(), value.format()])?; } + + // Format any end-of-line comments. + let mut first = true; + for range in stmt.trivia.iter().filter_map(|trivia| { + if matches!(trivia.relationship, Relationship::Trailing) { + if let TriviaKind::EndOfLineComment(range) = trivia.kind { + Some(range) + } else { + None + } + } else { + None + } + }) { + if std::mem::take(&mut first) { + write!(f, [line_suffix(&text(" "))])?; + } + write!(f, [line_suffix(&literal(range))])?; + } + + Ok(()) +} + +fn format_try( + f: &mut Formatter>, + stmt: &Stmt, + body: &[Stmt], + handlers: &[Excepthandler], + orelse: &[Stmt], + finalbody: &[Stmt], +) -> FormatResult<()> { + write!(f, [text("try:"), block_indent(&block(body))])?; + for handler in handlers { + write!(f, [handler.format()])?; + } + if !orelse.is_empty() { + write!(f, [text("else:"), block_indent(&block(orelse))])?; + } + if !finalbody.is_empty() { + write!(f, [text("finally:"), block_indent(&block(finalbody))])?; + } Ok(()) } @@ -720,7 +764,7 @@ impl Format> for FormatStmt<'_> { body, decorator_list, } => format_class_def(f, name, bases, keywords, body, decorator_list), - StmtKind::Return { value } => format_return(f, value.as_ref()), + StmtKind::Return { value } => format_return(f, self.item, value.as_ref()), StmtKind::Delete { targets } => format_delete(f, targets), StmtKind::Assign { targets, value, .. } => format_assign(f, self.item, targets, value), // StmtKind::AugAssign { .. } => {} @@ -778,7 +822,12 @@ impl Format> for FormatStmt<'_> { StmtKind::Raise { exc, cause } => { format_raise(f, self.item, exc.as_deref(), cause.as_deref()) } - // StmtKind::Try { .. } => {} + StmtKind::Try { + body, + handlers, + orelse, + finalbody, + } => format_try(f, self.item, body, handlers, orelse, finalbody), StmtKind::Assert { test, msg } => { format_assert(f, self.item, test, msg.as_ref().map(|expr| &**expr)) } diff --git a/crates/ruff_python_formatter/src/lib.rs b/crates/ruff_python_formatter/src/lib.rs index 0ac6633505..456b6c834e 100644 --- a/crates/ruff_python_formatter/src/lib.rs +++ b/crates/ruff_python_formatter/src/lib.rs @@ -102,6 +102,9 @@ mod tests { #[test_case(Path::new("simple_cases/tupleassign.py"); "tupleassign")] // Lots of deviations, _mostly_ related to string normalization and wrapping. #[test_case(Path::new("simple_cases/expression.py"); "expression")] + // Passing apart from a trailing end-of-line comment within an if statement, which is being + // inappropriately associated with the if statement rather than the line it's on. + #[test_case(Path::new("simple_cases/comments.py"); "comments")] #[test_case(Path::new("simple_cases/function.py"); "function")] #[test_case(Path::new("simple_cases/function2.py"); "function2")] #[test_case(Path::new("simple_cases/function_trailing_comma.py"); "function_trailing_comma")] diff --git a/crates/ruff_python_formatter/src/newlines.rs b/crates/ruff_python_formatter/src/newlines.rs index 33f7ced3ba..0e42e5a73a 100644 --- a/crates/ruff_python_formatter/src/newlines.rs +++ b/crates/ruff_python_formatter/src/newlines.rs @@ -48,7 +48,10 @@ impl<'a> Visitor<'a> for NewlineNormalizer { // Remove any runs of empty lines greater than two in a row. let mut count = 0; stmt.trivia.retain(|c| { - if matches!(c.kind, TriviaKind::EmptyLine) { + if matches!( + (c.kind, c.relationship), + (TriviaKind::EmptyLine, Relationship::Leading) + ) { count += 1; count <= self.depth.max_newlines() } else { @@ -64,7 +67,10 @@ impl<'a> Visitor<'a> for NewlineNormalizer { if seen_non_empty { true } else { - if matches!(c.kind, TriviaKind::EmptyLine) { + if matches!( + (c.kind, c.relationship), + (TriviaKind::EmptyLine, Relationship::Leading) + ) { false } else { seen_non_empty = true; @@ -87,7 +93,12 @@ impl<'a> Visitor<'a> for NewlineNormalizer { let present_newlines = stmt .trivia .iter() - .take_while(|c| matches!(c.kind, TriviaKind::EmptyLine)) + .take_while(|c| { + matches!( + (c.kind, c.relationship), + (TriviaKind::EmptyLine, Relationship::Leading) + ) + }) .count(); if present_newlines < required_newlines { for _ in 0..(required_newlines - present_newlines) { @@ -113,7 +124,12 @@ impl<'a> Visitor<'a> for NewlineNormalizer { - stmt .trivia .iter() - .take_while(|c| matches!(c.kind, TriviaKind::EmptyLine)) + .take_while(|c| { + matches!( + (c.kind, c.relationship), + (TriviaKind::EmptyLine, Relationship::Leading) + ) + }) .count(); for _ in 0..num_to_insert { stmt.trivia.insert( diff --git a/crates/ruff_python_formatter/src/trivia.rs b/crates/ruff_python_formatter/src/trivia.rs index 4891082d01..dd721faed0 100644 --- a/crates/ruff_python_formatter/src/trivia.rs +++ b/crates/ruff_python_formatter/src/trivia.rs @@ -128,7 +128,7 @@ pub fn extract_trivia_tokens(lxr: &[LexResult]) -> Vec { let mut parens = vec![]; for (start, tok, end) in lxr.iter().flatten() { // Add empty lines. - if let Some((prev, ..)) = prev_non_newline_tok { + if let Some((.., prev)) = prev_non_newline_tok { for row in prev.row() + 1..start.row() { tokens.push(TriviaToken { start: Location::new(row, 0), @@ -189,7 +189,7 @@ pub fn extract_trivia_tokens(lxr: &[LexResult]) -> Vec { prev_tok = Some((start, tok, end)); // Track the most recent non-whitespace token. - if !matches!(tok, Tok::Newline | Tok::NonLogicalNewline,) { + if !matches!(tok, Tok::Newline | Tok::NonLogicalNewline) { prev_non_newline_tok = Some((start, tok, end)); }