From d6daa615637aa47dcaca34f051eef8e6d2bdda6c Mon Sep 17 00:00:00 2001 From: Micha Reiser Date: Sat, 3 Jun 2023 15:29:22 +0200 Subject: [PATCH] Handle trailing end-of-line comments in-between-bodies (#4812) ## Summary And more custom logic around comments in bodies... uff. Let's say we have the following code ```python if x == y: pass # trailing comment of pass else: # trailing comment of `else` print("I have no comments") ``` Right now, the formatter attaches the `# trailing comment of `else` as a trailing comment of `pass` because it doesn't "see" that there's an `else` keyword in between (because the else body is just a Vec and not a node). This PR adds custom logic that attaches the trailing comments after the `else` as dangling comments to the `if` statement. The if statement must then split the dangling comments by `comments.text_position()`: * All comments up to the first end-of-line comment are leading comments of the `else` keyword. * All end-of-line comments coming after are `trailing` comments for the `else` keyword. ## Test Plan I added new unit tests. --- .../ruff_python_formatter/src/comments/mod.rs | 15 ++ .../src/comments/placement.rs | 155 +++++++++++++----- ...ile_trailing_else_end_of_line_comment.snap | 21 +++ 3 files changed, 150 insertions(+), 41 deletions(-) create mode 100644 crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__while_trailing_else_end_of_line_comment.snap diff --git a/crates/ruff_python_formatter/src/comments/mod.rs b/crates/ruff_python_formatter/src/comments/mod.rs index 6792adfcd2..57206e3fc4 100644 --- a/crates/ruff_python_formatter/src/comments/mod.rs +++ b/crates/ruff_python_formatter/src/comments/mod.rs @@ -872,4 +872,19 @@ a = ( assert_debug_snapshot!(comments.debug(test_case.source_code)); } + + #[test] + fn while_trailing_else_end_of_line_comment() { + let source = r#"while True: + pass +else: # trailing comment + pass +"#; + + let test_case = CommentsTestCase::from_code(source); + + let comments = test_case.to_comments(); + + assert_debug_snapshot!(comments.debug(test_case.source_code)); + } } diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs index 23bc961c5f..d4b007e775 100644 --- a/crates/ruff_python_formatter/src/comments/placement.rs +++ b/crates/ruff_python_formatter/src/comments/placement.rs @@ -16,7 +16,8 @@ pub(super) fn place_comment<'a>( ) -> CommentPlacement<'a> { handle_in_between_excepthandlers_or_except_handler_and_else_or_finally_comment(comment, locator) .or_else(|comment| handle_match_comment(comment, locator)) - .or_else(|comment| handle_in_between_bodies_comment(comment, locator)) + .or_else(|comment| handle_in_between_bodies_own_line_comment(comment, locator)) + .or_else(|comment| handle_in_between_bodies_end_of_line_comment(comment, locator)) .or_else(|comment| handle_trailing_body_comment(comment, locator)) .or_else(handle_trailing_end_of_line_body_comment) .or_else(|comment| handle_positional_only_arguments_separator_comment(comment, locator)) @@ -178,7 +179,7 @@ fn handle_in_between_excepthandlers_or_except_handler_and_else_or_finally_commen CommentPlacement::Default(comment) } -/// Handles comments between the last statement and the first statement of two bodies. +/// Handles own line comments between the last statement and the first statement of two bodies. /// /// ```python /// if x == y: @@ -188,15 +189,11 @@ fn handle_in_between_excepthandlers_or_except_handler_and_else_or_finally_commen /// else: /// print("I have no comments") /// ``` -fn handle_in_between_bodies_comment<'a>( +fn handle_in_between_bodies_own_line_comment<'a>( comment: DecoratedComment<'a>, locator: &Locator, ) -> CommentPlacement<'a> { - use ruff_python_ast::prelude::*; - - // The rule only applies to own line comments. The default logic associates end of line comments - // correctly. - if comment.text_position().is_end_of_line() { + if !comment.text_position().is_own_line() { return CommentPlacement::Default(comment); } @@ -204,39 +201,7 @@ fn handle_in_between_bodies_comment<'a>( if let (Some(preceding), Some(following)) = (comment.preceding_node(), comment.following_node()) { // ...and the following statement must be the first statement in an alternate body of the parent... - let is_following_the_first_statement_in_a_parents_alternate_body = - match comment.enclosing_node() { - AnyNodeRef::StmtIf(StmtIf { orelse, .. }) - | AnyNodeRef::StmtFor(StmtFor { orelse, .. }) - | AnyNodeRef::StmtAsyncFor(StmtAsyncFor { orelse, .. }) - | AnyNodeRef::StmtWhile(StmtWhile { orelse, .. }) => { - are_same_optional(following, orelse.first()) - } - - AnyNodeRef::StmtTry(StmtTry { - handlers, - orelse, - finalbody, - .. - }) - | AnyNodeRef::StmtTryStar(StmtTryStar { - handlers, - orelse, - finalbody, - .. - }) => { - are_same_optional(following, handlers.first()) - // Comments between the handlers and the `else`, or comments between the `handlers` and the `finally` - // are already handled by `handle_in_between_excepthandlers_or_except_handler_and_else_or_finally_comment` - || handlers.is_empty() && are_same_optional(following, orelse.first()) - || (handlers.is_empty() || !orelse.is_empty()) - && are_same_optional(following, finalbody.first()) - } - - _ => false, - }; - - if !is_following_the_first_statement_in_a_parents_alternate_body { + if !is_first_statement_in_enclosing_alternate_body(following, comment.enclosing_node()) { // ```python // if test: // a @@ -305,6 +270,75 @@ fn handle_in_between_bodies_comment<'a>( CommentPlacement::Default(comment) } +/// Handles end of line comments comments between the last statement and the first statement of two bodies. +/// +/// ```python +/// if x == y: +/// pass # trailing comment of pass +/// else: # trailing comment of `else` +/// print("I have no comments") +/// ``` +fn handle_in_between_bodies_end_of_line_comment<'a>( + comment: DecoratedComment<'a>, + locator: &Locator, +) -> CommentPlacement<'a> { + if !comment.text_position().is_end_of_line() { + return CommentPlacement::Default(comment); + } + + // The comment must be between two statements... + if let (Some(preceding), Some(following)) = (comment.preceding_node(), comment.following_node()) + { + // ...and the following statement must be the first statement in an alternate body of the parent... + if !is_first_statement_in_enclosing_alternate_body(following, comment.enclosing_node()) { + // ```python + // if test: + // a + // # comment + // b + // ``` + return CommentPlacement::Default(comment); + } + + if !locator.contains_line_break(TextRange::new(preceding.end(), comment.slice().start())) { + // Trailing comment of the preceding statement + // ```python + // while test: + // a # comment + // else: + // b + // ``` + CommentPlacement::trailing(preceding, comment) + } else if following.is_stmt_if() || following.is_except_handler() { + // The `elif` or except handlers have their own body to which we can attach the trailing comment + // ```python + // if test: + // a + // elif c: # comment + // b + // ``` + CommentPlacement::trailing(following, comment) + } else { + // There are no bodies for the "else" branch and other bodies that are represented as a `Vec`. + // This means, there's no good place to attach the comments to. + // Make this a dangling comments and manually format the comment in + // in the enclosing node's formatting logic. For `try`, it's the formatters responsibility + // to correctly identify the comments for the `finally` and `orelse` block by looking + // at the comment's range. + // + // ```python + // while x == y: + // pass + // else: # trailing + // print("nooop") + // ``` + CommentPlacement::dangling(comment.enclosing_node(), comment) + } + } else { + CommentPlacement::Default(comment) + } +} + /// Handles trailing comments at the end of a body block (or any other block that is indented). /// ```python /// def test(): @@ -703,3 +737,42 @@ fn last_child_in_body(node: AnyNodeRef) -> Option { body.last().map(AnyNodeRef::from) } + +/// Returns `true` if `following` is the first statement in an alternate `body` (e.g. the else of an if statement) of the `enclosing` node. +fn is_first_statement_in_enclosing_alternate_body( + following: AnyNodeRef, + enclosing: AnyNodeRef, +) -> bool { + use ruff_python_ast::prelude::*; + + match enclosing { + AnyNodeRef::StmtIf(StmtIf { orelse, .. }) + | AnyNodeRef::StmtFor(StmtFor { orelse, .. }) + | AnyNodeRef::StmtAsyncFor(StmtAsyncFor { orelse, .. }) + | AnyNodeRef::StmtWhile(StmtWhile { orelse, .. }) => { + are_same_optional(following, orelse.first()) + } + + AnyNodeRef::StmtTry(StmtTry { + handlers, + orelse, + finalbody, + .. + }) + | AnyNodeRef::StmtTryStar(StmtTryStar { + handlers, + orelse, + finalbody, + .. + }) => { + are_same_optional(following, handlers.first()) + // Comments between the handlers and the `else`, or comments between the `handlers` and the `finally` + // are already handled by `handle_in_between_excepthandlers_or_except_handler_and_else_or_finally_comment` + || handlers.is_empty() && are_same_optional(following, orelse.first()) + || (handlers.is_empty() || !orelse.is_empty()) + && are_same_optional(following, finalbody.first()) + } + + _ => false, + } +} diff --git a/crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__while_trailing_else_end_of_line_comment.snap b/crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__while_trailing_else_end_of_line_comment.snap new file mode 100644 index 0000000000..33cd36edc4 --- /dev/null +++ b/crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__while_trailing_else_end_of_line_comment.snap @@ -0,0 +1,21 @@ +--- +source: crates/ruff_python_formatter/src/comments/mod.rs +expression: comments.debug(test_case.source_code) +--- +{ + Node { + kind: StmtWhile, + range: 0..54, + source: `while True:⏎`, + }: { + "leading": [], + "dangling": [ + SourceComment { + text: "# trailing comment", + position: EndOfLine, + formatted: false, + }, + ], + "trailing": [], + }, +}