Add StmtKind::Try; fix trailing newlines (#3074)

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<ASTFormatContext<'_>> for Excepthandler {
+    type Format<'a> = FormatExcepthandler<'a>;
+
+    fn format(&self) -> Self::Format<'_> {
+        FormatExcepthandler { item: self }
+    }
+}
+
+impl Format<ASTFormatContext<'_>> for FormatExcepthandler<'_> {
+    fn fmt(&self, f: &mut Formatter<ASTFormatContext>) -> 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(())
+    }
+}

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;

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<ASTFormatContext<'_>>,
+    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<ASTFormatContext<'_>>,
+    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<ASTFormatContext<'_>> 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<ASTFormatContext<'_>> 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))
             }

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")]

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(

crates/ruff_python_formatter/src/trivia.rs

@@ -128,7 +128,7 @@ pub fn extract_trivia_tokens(lxr: &[LexResult]) -> Vec<TriviaToken> {
     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<TriviaToken> {
         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));
         }