diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs index a84c2e3237..7653fa4fc2 100644 --- a/crates/ruff_python_formatter/src/comments/placement.rs +++ b/crates/ruff_python_formatter/src/comments/placement.rs @@ -25,24 +25,14 @@ pub(super) fn place_comment<'a>( ) -> CommentPlacement<'a> { // Handle comments before and after bodies such as the different branches of an if statement. let comment = if comment.line_position().is_own_line() { - match handle_own_line_comment_around_body(comment, locator) { - CommentPlacement::Default(comment) => comment, - placement => { - return placement; - } - } + handle_own_line_comment_around_body(comment, locator) } else { - match handle_end_of_line_comment_around_body(comment, locator) { - CommentPlacement::Default(comment) => comment, - placement => { - return placement; - } - } + handle_end_of_line_comment_around_body(comment, locator) }; // Change comment placement depending on the node type. These can be seen as node-specific // fixups. - match comment.enclosing_node() { + comment.then_with(|comment| match comment.enclosing_node() { AnyNodeRef::Parameters(arguments) => { handle_parameters_separator_comment(comment, arguments, locator) } @@ -88,7 +78,7 @@ pub(super) fn place_comment<'a>( } AnyNodeRef::StmtImportFrom(import_from) => handle_import_from_comment(comment, import_from), _ => CommentPlacement::Default(comment), - } + }) } fn handle_end_of_line_comment_around_body<'a>( @@ -275,22 +265,18 @@ fn handle_own_line_comment_around_body<'a>( } // Check if we're between bodies and should attach to the following body. - let comment = match handle_own_line_comment_between_branches(comment, preceding, locator) { - CommentPlacement::Default(comment) => comment, - placement => return placement, - }; - - // Otherwise, there's no following branch or the indentation is too deep, so attach to the - // recursively last statement in the preceding body with the matching indentation. - let comment = match handle_own_line_comment_after_branch(comment, preceding, locator) { - CommentPlacement::Default(comment) => comment, - placement => return placement, - }; - - // If the following node is the first in its body, and there's a non-trivia token between the - // comment and the following node (like a parenthesis), then it means the comment is trailing - // the preceding node, not leading the following one. - handle_own_line_comment_in_clause(comment, preceding, locator) + handle_own_line_comment_between_branches(comment, preceding, locator) + .then_with(|comment| { + // Otherwise, there's no following branch or the indentation is too deep, so attach to the + // recursively last statement in the preceding body with the matching indentation. + handle_own_line_comment_after_branch(comment, preceding, locator) + }) + .then_with(|comment| { + // If the following node is the first in its body, and there's a non-trivia token between the + // comment and the following node (like a parenthesis), then it means the comment is trailing + // the preceding node, not leading the following one. + handle_own_line_comment_in_clause(comment, preceding, locator) + }) } /// Handles own line comments between two branches of a node. diff --git a/crates/ruff_python_formatter/src/comments/visitor.rs b/crates/ruff_python_formatter/src/comments/visitor.rs index 8f2d61f016..467e777f13 100644 --- a/crates/ruff_python_formatter/src/comments/visitor.rs +++ b/crates/ruff_python_formatter/src/comments/visitor.rs @@ -658,6 +658,17 @@ impl<'a> CommentPlacement<'a> { comment: comment.into(), } } + + /// Chains the placement with the given function. + /// + /// Returns `self` when the placement is non-[`CommentPlacement::Default`]. Otherwise, calls the + /// function with the comment and returns the result. + pub(super) fn then_with) -> Self>(self, f: F) -> Self { + match self { + Self::Default(comment) => f(comment), + _ => self, + } + } } #[derive(Copy, Clone, Eq, PartialEq, Debug)]