diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/try.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/try.py index f50dcab4c4..11d07349bc 100644 --- a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/try.py +++ b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/try.py @@ -70,3 +70,22 @@ try: except: a = 10 # trailing comment1 b = 11 # trailing comment2 + + +# try/except*, mostly the same as try +try: # try + ... + # end of body +# before except +except* (Exception, ValueError) as exc: # except line + ... +# before except 2 +except* KeyError as key: # except line 2 + ... + # in body 2 +# before else +else: + ... +# before finally +finally: + ... 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 2e1c695217..66ab84be59 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 @@ -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> + for FormatExceptHandlerExceptHandler +{ + type Options = ExceptHandlerKind; + + fn with_options(mut self, options: Self::Options) -> Self { + self.except_handler_kind = options; + self + } +} impl FormatNodeRule for FormatExceptHandlerExceptHandler { fn fmt_fields( @@ -25,7 +46,16 @@ impl FormatNodeRule 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!( diff --git a/crates/ruff_python_formatter/src/statement/stmt_try.rs b/crates/ruff_python_formatter/src/statement/stmt_try.rs index 701fb8165b..49bcef4387 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_try.rs +++ b/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> for FormatExceptHandler { + type Options = ExceptHandlerKind; + + fn with_options(mut self, options: Self::Options) -> Self { + self.except_handler_kind = options; + self + } +} impl FormatRule> for FormatExceptHandler { fn fmt( @@ -22,7 +106,9 @@ impl FormatRule> for FormatExceptHandler { f: &mut Formatter>, ) -> 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> for ExceptHandler { FormatRefWithRule::new(self, FormatExceptHandler::default()) } } - -impl FormatNodeRule for FormatStmtTry { - fn fmt_fields(&self, item: &StmtTry, f: &mut PyFormatter) -> FormatResult<()> { - let StmtTry { - range: _, - body, - handlers, - orelse, - finalbody, - } = item; - +impl Format> for AnyStatementTry<'_> { + fn fmt(&self, f: &mut Formatter>) -> 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 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 for FormatStmtTry { write!(f, [comments::dangling_comments(dangling_comments)]) } +} + +impl FormatNodeRule 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(()) } } diff --git a/crates/ruff_python_formatter/src/statement/stmt_try_star.rs b/crates/ruff_python_formatter/src/statement/stmt_try_star.rs index aa4d45c444..3c61258248 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_try_star.rs +++ b/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 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(()) } } diff --git a/crates/ruff_python_formatter/tests/snapshots/format@statement__try.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@statement__try.py.snap index e0a6fa67b0..588c6e5b50 100644 --- a/crates/ruff_python_formatter/tests/snapshots/format@statement__try.py.snap +++ b/crates/ruff_python_formatter/tests/snapshots/format@statement__try.py.snap @@ -76,6 +76,25 @@ try: except: a = 10 # trailing comment1 b = 11 # trailing comment2 + + +# try/except*, mostly the same as try +try: # try + ... + # end of body +# before except +except* (Exception, ValueError) as exc: # except line + ... +# before except 2 +except* KeyError as key: # except line 2 + ... + # in body 2 +# before else +else: + ... +# before finally +finally: + ... ``` ## Output @@ -161,6 +180,26 @@ try: except: a = 10 # trailing comment1 b = 11 # trailing comment2 + + +# try/except*, mostly the same as try +try: + # try + ... + # end of body +# before except +except* (Exception, ValueError) as exc: # except line + ... +# before except 2 +except* KeyError as key: # except line 2 + ... + # in body 2 +# before else +else: + ... +# before finally +finally: + ... ```