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))
             }