format StmtTryStar (#5418)

crates/ruff_python_formatter/src/other/except_handler_except_handler.rs

@@ -2,12 +2,33 @@ use crate::comments::trailing_comments;
 use crate::expression::parentheses::Parenthesize;
 use crate::prelude::*;
 use crate::{FormatNodeRule, PyFormatter};
+use ruff_formatter::FormatRuleWithOptions;
 use ruff_formatter::{write, Buffer, FormatResult};
 use ruff_python_ast::node::AstNode;
 use rustpython_parser::ast::ExceptHandlerExceptHandler;
 
+#[derive(Copy, Clone, Default)]
+pub enum ExceptHandlerKind {
+    #[default]
+    Regular,
+    Starred,
+}
+
 #[derive(Default)]
-pub struct FormatExceptHandlerExceptHandler;
+pub struct FormatExceptHandlerExceptHandler {
+    except_handler_kind: ExceptHandlerKind,
+}
+
+impl FormatRuleWithOptions<ExceptHandlerExceptHandler, PyFormatContext<'_>>
+    for FormatExceptHandlerExceptHandler
+{
+    type Options = ExceptHandlerKind;
+
+    fn with_options(mut self, options: Self::Options) -> Self {
+        self.except_handler_kind = options;
+        self
+    }
+}
 
 impl FormatNodeRule<ExceptHandlerExceptHandler> for FormatExceptHandlerExceptHandler {
     fn fmt_fields(
@@ -25,7 +46,16 @@ impl FormatNodeRule<ExceptHandlerExceptHandler> for FormatExceptHandlerExceptHan
         let comments_info = f.context().comments().clone();
         let dangling_comments = comments_info.dangling_comments(item.as_any_node_ref());
 
-        write!(f, [text("except")])?;
+        write!(
+            f,
+            [
+                text("except"),
+                match self.except_handler_kind {
+                    ExceptHandlerKind::Regular => None,
+                    ExceptHandlerKind::Starred => Some(text("*")),
+                }
+            ]
+        )?;
 
         if let Some(type_) = type_ {
             write!(

crates/ruff_python_formatter/src/statement/stmt_try.rs

@@ -1,19 +1,103 @@
 use crate::comments;
 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::FormatRefWithRule;
 use crate::statement::Stmt;
 use crate::{FormatNodeRule, PyFormatter};
+use ruff_formatter::FormatRuleWithOptions;
 use ruff_formatter::{write, Buffer, FormatResult};
-use ruff_python_ast::node::AstNode;
-use rustpython_parser::ast::{ExceptHandler, Ranged, StmtTry, Suite};
+use ruff_python_ast::node::AnyNodeRef;
+use ruff_text_size::TextRange;
+use rustpython_parser::ast::{ExceptHandler, Ranged, StmtTry, StmtTryStar, Suite};
+
+pub(super) enum AnyStatementTry<'a> {
+    Try(&'a StmtTry),
+    TryStar(&'a StmtTryStar),
+}
+impl<'a> AnyStatementTry<'a> {
+    const fn except_handler_kind(&self) -> ExceptHandlerKind {
+        match self {
+            AnyStatementTry::Try(_) => ExceptHandlerKind::Regular,
+            AnyStatementTry::TryStar(_) => ExceptHandlerKind::Starred,
+        }
+    }
+
+    fn body(&self) -> &Suite {
+        match self {
+            AnyStatementTry::Try(try_) => &try_.body,
+            AnyStatementTry::TryStar(try_) => &try_.body,
+        }
+    }
+
+    fn handlers(&self) -> &[ExceptHandler] {
+        match self {
+            AnyStatementTry::Try(try_) => try_.handlers.as_slice(),
+            AnyStatementTry::TryStar(try_) => try_.handlers.as_slice(),
+        }
+    }
+    fn orelse(&self) -> &Suite {
+        match self {
+            AnyStatementTry::Try(try_) => &try_.orelse,
+            AnyStatementTry::TryStar(try_) => &try_.orelse,
+        }
+    }
+
+    fn finalbody(&self) -> &Suite {
+        match self {
+            AnyStatementTry::Try(try_) => &try_.finalbody,
+            AnyStatementTry::TryStar(try_) => &try_.finalbody,
+        }
+    }
+}
+
+impl Ranged for AnyStatementTry<'_> {
+    fn range(&self) -> TextRange {
+        match self {
+            AnyStatementTry::Try(with) => with.range(),
+            AnyStatementTry::TryStar(with) => with.range(),
+        }
+    }
+}
+
+impl<'a> From<&'a StmtTry> for AnyStatementTry<'a> {
+    fn from(value: &'a StmtTry) -> Self {
+        AnyStatementTry::Try(value)
+    }
+}
+
+impl<'a> From<&'a StmtTryStar> for AnyStatementTry<'a> {
+    fn from(value: &'a StmtTryStar) -> Self {
+        AnyStatementTry::TryStar(value)
+    }
+}
+
+impl<'a> From<&AnyStatementTry<'a>> for AnyNodeRef<'a> {
+    fn from(value: &AnyStatementTry<'a>) -> Self {
+        match value {
+            AnyStatementTry::Try(with) => AnyNodeRef::StmtTry(with),
+            AnyStatementTry::TryStar(with) => AnyNodeRef::StmtTryStar(with),
+        }
+    }
+}
 
 #[derive(Default)]
 pub struct FormatStmtTry;
 
 #[derive(Copy, Clone, Default)]
-pub struct FormatExceptHandler;
+pub struct FormatExceptHandler {
+    except_handler_kind: ExceptHandlerKind,
+}
+
+impl FormatRuleWithOptions<ExceptHandler, PyFormatContext<'_>> for FormatExceptHandler {
+    type Options = ExceptHandlerKind;
+
+    fn with_options(mut self, options: Self::Options) -> Self {
+        self.except_handler_kind = options;
+        self
+    }
+}
 
 impl FormatRule<ExceptHandler, PyFormatContext<'_>> for FormatExceptHandler {
     fn fmt(
@@ -22,7 +106,9 @@ impl FormatRule<ExceptHandler, PyFormatContext<'_>> for FormatExceptHandler {
         f: &mut Formatter<PyFormatContext<'_>>,
     ) -> FormatResult<()> {
         match item {
-            ExceptHandler::ExceptHandler(x) => x.format().fmt(f),
+            ExceptHandler::ExceptHandler(x) => {
+                x.format().with_options(self.except_handler_kind).fmt(f)
+            }
         }
     }
 }
@@ -39,19 +125,14 @@ impl<'ast> AsFormat<PyFormatContext<'ast>> for ExceptHandler {
         FormatRefWithRule::new(self, FormatExceptHandler::default())
     }
 }
-
-impl FormatNodeRule<StmtTry> for FormatStmtTry {
-    fn fmt_fields(&self, item: &StmtTry, f: &mut PyFormatter) -> FormatResult<()> {
-        let StmtTry {
-            range: _,
-            body,
-            handlers,
-            orelse,
-            finalbody,
-        } = item;
-
+impl Format<PyFormatContext<'_>> for AnyStatementTry<'_> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
         let comments_info = f.context().comments().clone();
-        let mut dangling_comments = comments_info.dangling_comments(item.as_any_node_ref());
+        let mut dangling_comments = comments_info.dangling_comments(self);
+        let body = self.body();
+        let handlers = self.handlers();
+        let orelse = self.orelse();
+        let finalbody = self.finalbody();
 
         write!(f, [text("try:"), block_indent(&body.format())])?;
 
@@ -63,7 +144,7 @@ impl FormatNodeRule<StmtTry> for FormatStmtTry {
                 f,
                 [
                     leading_alternate_branch_comments(handler_comments, previous_node),
-                    &handler.format()
+                    &handler.format().with_options(self.except_handler_kind()),
                 ]
             )?;
             previous_node = match handler {
@@ -78,9 +159,15 @@ impl FormatNodeRule<StmtTry> for FormatStmtTry {
 
         write!(f, [comments::dangling_comments(dangling_comments)])
     }
+}
+
+impl FormatNodeRule<StmtTry> for FormatStmtTry {
+    fn fmt_fields(&self, item: &StmtTry, f: &mut PyFormatter) -> FormatResult<()> {
+        AnyStatementTry::from(item).fmt(f)
+    }
 
     fn fmt_dangling_comments(&self, _node: &StmtTry, _f: &mut PyFormatter) -> FormatResult<()> {
-        // dangling comments are formatted as part of fmt_fields
+        // dangling comments are formatted as part of AnyStatementTry::fmt
         Ok(())
     }
 }

crates/ruff_python_formatter/src/statement/stmt_try_star.rs

@@ -1,5 +1,7 @@
-use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
+use crate::statement::stmt_try::AnyStatementTry;
+use crate::{FormatNodeRule, PyFormatter};
+use ruff_formatter::Format;
+use ruff_formatter::FormatResult;
 use rustpython_parser::ast::StmtTryStar;
 
 #[derive(Default)]
@@ -7,6 +9,11 @@ pub struct FormatStmtTryStar;
 
 impl FormatNodeRule<StmtTryStar> for FormatStmtTryStar {
     fn fmt_fields(&self, item: &StmtTryStar, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [not_yet_implemented(item)])
+        AnyStatementTry::from(item).fmt(f)
+    }
+
+    fn fmt_dangling_comments(&self, _node: &StmtTryStar, _f: &mut PyFormatter) -> FormatResult<()> {
+        // dangling comments are formatted as part of AnyStatementTry::fmt
+        Ok(())
     }
 }