From c3a9151eb5758ae1647d1a727ef87492bfc71c74 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 14 Aug 2023 11:11:03 -0400 Subject: [PATCH] Handle comments on open parentheses in with statements (#6515) ## Summary This PR adds handling for comments on open parentheses in parenthesized context managers. For example, given: ```python with ( # comment CtxManager1() as example1, CtxManager2() as example2, CtxManager3() as example3, ): ... ``` We want to preserve that formatting. (Black does the same.) On `main`, we format as: ```python with ( # comment CtxManager1() as example1, CtxManager2() as example2, CtxManager3() as example3, ): ... ``` It's very similar to how `StmtImportFrom` is handled. Note that this case _isn't_ covered by the "parenthesized comment" proposal, since this is a common on the statement that would typically be attached to the first `WithItem`, and the `WithItem` _itself_ can have parenthesized comments, like: ```python with ( # comment ( CtxManager1() # comment ) as example1, CtxManager2() as example2, CtxManager3() as example3, ): ... ``` ## Test Plan `cargo test` Confirmed no change in similarity score. --- .../test/fixtures/ruff/statement/with.py | 18 ++++++++ .../src/comments/placement.rs | 36 ++++++++++++++-- .../src/statement/stmt_import_from.rs | 27 ++++++++---- .../src/statement/stmt_with.rs | 42 ++++++++++++++++--- ..._opening_parentheses_comment_value.py.snap | 3 +- .../format@statement__import_from.py.snap | 4 +- .../snapshots/format@statement__with.py.snap | 37 ++++++++++++++++ 7 files changed, 148 insertions(+), 19 deletions(-) diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/with.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/with.py index 76a8a49950..8a231ed788 100644 --- a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/with.py +++ b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/with.py @@ -116,3 +116,21 @@ with [ dddddddddddddddddddddddddddddddd, ] as example1, aaaaaaaaaaaaaaaaaaaaaaaaaa * bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb * cccccccccccccccccccccccccccc + ddddddddddddddddd as example2, CtxManager222222222222222() as example2: ... + +# Comments on open parentheses +with ( # comment + CtxManager1() as example1, + CtxManager2() as example2, + CtxManager3() as example3, +): + ... + + +with ( # outer comment + ( # inner comment + CtxManager1() + ) as example1, + CtxManager2() as example2, + CtxManager3() as example3, +): + ... diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs index d8d8c03578..685aa2d3b8 100644 --- a/crates/ruff_python_formatter/src/comments/placement.rs +++ b/crates/ruff_python_formatter/src/comments/placement.rs @@ -73,6 +73,7 @@ pub(super) fn place_comment<'a>( handle_leading_class_with_decorators_comment(comment, class_def) } AnyNodeRef::StmtImportFrom(import_from) => handle_import_from_comment(comment, import_from), + AnyNodeRef::StmtWith(with_) => handle_with_comment(comment, with_), AnyNodeRef::ExprConstant(_) => { if let Some(AnyNodeRef::ExprFString(fstring)) = comment.enclosing_parent() { CommentPlacement::dangling(fstring, comment) @@ -1172,7 +1173,7 @@ fn handle_bracketed_end_of_line_comment<'a>( CommentPlacement::Default(comment) } -/// Attach an enclosed end-of-line comment to a [`StmtImportFrom`]. +/// Attach an enclosed end-of-line comment to a [`ast::StmtImportFrom`]. /// /// For example, given: /// ```python @@ -1181,8 +1182,8 @@ fn handle_bracketed_end_of_line_comment<'a>( /// ) /// ``` /// -/// The comment will be attached to the `StmtImportFrom` node as a dangling comment, to ensure -/// that it remains on the same line as the `StmtImportFrom` itself. +/// The comment will be attached to the [`ast::StmtImportFrom`] node as a dangling comment, to +/// ensure that it remains on the same line as the [`ast::StmtImportFrom`] itself. fn handle_import_from_comment<'a>( comment: DecoratedComment<'a>, import_from: &'a ast::StmtImportFrom, @@ -1217,6 +1218,35 @@ fn handle_import_from_comment<'a>( } } +/// Attach an enclosed end-of-line comment to a [`ast::StmtWith`]. +/// +/// For example, given: +/// ```python +/// with ( # foo +/// CtxManager1() as example1, +/// CtxManager2() as example2, +/// CtxManager3() as example3, +/// ): +/// ... +/// ``` +/// +/// The comment will be attached to the [`ast::StmtWith`] node as a dangling comment, to ensure +/// that it remains on the same line as the [`ast::StmtWith`] itself. +fn handle_with_comment<'a>( + comment: DecoratedComment<'a>, + with_statement: &'a ast::StmtWith, +) -> CommentPlacement<'a> { + if comment.line_position().is_end_of_line() + && with_statement.items.first().is_some_and(|with_item| { + with_statement.start() < comment.start() && comment.start() < with_item.start() + }) + { + CommentPlacement::dangling(comment.enclosing_node(), comment) + } else { + CommentPlacement::Default(comment) + } +} + // Handle comments inside comprehensions, e.g. // // ```python diff --git a/crates/ruff_python_formatter/src/statement/stmt_import_from.rs b/crates/ruff_python_formatter/src/statement/stmt_import_from.rs index 30b76d57b6..3249dd7a79 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_import_from.rs +++ b/crates/ruff_python_formatter/src/statement/stmt_import_from.rs @@ -4,7 +4,7 @@ use ruff_python_ast::node::AstNode; use ruff_python_ast::{Ranged, StmtImportFrom}; use crate::builders::{parenthesize_if_expands, PyFormatterExtensions}; -use crate::comments::trailing_comments; +use crate::expression::parentheses::parenthesized; use crate::{AsFormat, FormatNodeRule, PyFormatter}; #[derive(Default)] @@ -15,8 +15,8 @@ impl FormatNodeRule for FormatStmtImportFrom { let StmtImportFrom { module, names, - range: _, level, + range: _, } = item; let level_str = level @@ -43,16 +43,29 @@ impl FormatNodeRule for FormatStmtImportFrom { } } - let comments = f.context().comments().clone(); - let dangling_comments = comments.dangling_comments(item.as_any_node_ref()); - write!(f, [trailing_comments(dangling_comments)])?; - let names = format_with(|f| { f.join_comma_separated(item.end()) .entries(names.iter().map(|name| (name, name.format()))) .finish() }); - parenthesize_if_expands(&names).fmt(f) + + // A dangling comment on an import is a parenthesized comment, like: + // ```python + // from example import ( # comment + // A, + // B, + // ) + // ``` + let comments = f.context().comments().clone(); + let parenthesized_comments = comments.dangling_comments(item.as_any_node_ref()); + + if parenthesized_comments.is_empty() { + parenthesize_if_expands(&names).fmt(f) + } else { + parenthesized("(", &names, ")") + .with_dangling_comments(parenthesized_comments) + .fmt(f) + } } fn fmt_dangling_comments( diff --git a/crates/ruff_python_formatter/src/statement/stmt_with.rs b/crates/ruff_python_formatter/src/statement/stmt_with.rs index b8afb2042d..cb7493dbdd 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_with.rs +++ b/crates/ruff_python_formatter/src/statement/stmt_with.rs @@ -1,11 +1,12 @@ use ruff_formatter::{format_args, write, FormatError}; +use ruff_python_ast::node::AstNode; use ruff_python_ast::{Ranged, StmtWith}; use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer}; use ruff_text_size::TextRange; use crate::comments::trailing_comments; use crate::expression::parentheses::{ - in_parentheses_only_soft_line_break_or_space, optional_parentheses, + in_parentheses_only_soft_line_break_or_space, optional_parentheses, parenthesized, }; use crate::prelude::*; use crate::FormatNodeRule; @@ -15,9 +16,6 @@ pub struct FormatStmtWith; impl FormatNodeRule for FormatStmtWith { fn fmt_fields(&self, item: &StmtWith, f: &mut PyFormatter) -> FormatResult<()> { - let comments = f.context().comments().clone(); - let dangling_comments = comments.dangling_comments(item); - write!( f, [ @@ -28,7 +26,39 @@ impl FormatNodeRule for FormatStmtWith { ] )?; - if are_with_items_parenthesized(item, f.context())? { + // The `with` statement can have one dangling comment on the open parenthesis, like: + // ```python + // with ( # comment + // CtxManager() as example + // ): + // ... + // ``` + // + // Any other dangling comments are trailing comments on the colon, like: + // ```python + // with CtxManager() as example: # comment + // ... + // ``` + let comments = f.context().comments().clone(); + let dangling_comments = comments.dangling_comments(item.as_any_node_ref()); + let partition_point = dangling_comments.partition_point(|comment| { + item.items + .first() + .is_some_and(|with_item| with_item.start() > comment.start()) + }); + let (parenthesized_comments, colon_comments) = dangling_comments.split_at(partition_point); + + if !parenthesized_comments.is_empty() { + let joined = format_with(|f: &mut PyFormatter| { + f.join_comma_separated(item.body.first().unwrap().start()) + .nodes(&item.items) + .finish() + }); + + parenthesized("(", &joined, ")") + .with_dangling_comments(parenthesized_comments) + .fmt(f)?; + } else if are_with_items_parenthesized(item, f.context())? { optional_parentheses(&format_with(|f| { let mut joiner = f.join_comma_separated(item.body.first().unwrap().start()); @@ -52,7 +82,7 @@ impl FormatNodeRule for FormatStmtWith { f, [ text(":"), - trailing_comments(dangling_comments), + trailing_comments(colon_comments), block_indent(&item.body.format()) ] ) diff --git a/crates/ruff_python_formatter/tests/snapshots/format@parentheses__opening_parentheses_comment_value.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@parentheses__opening_parentheses_comment_value.py.snap index 9ba9c4892c..e4a42c950f 100644 --- a/crates/ruff_python_formatter/tests/snapshots/format@parentheses__opening_parentheses_comment_value.py.snap +++ b/crates/ruff_python_formatter/tests/snapshots/format@parentheses__opening_parentheses_comment_value.py.snap @@ -170,8 +170,7 @@ async def h(): ) -with ( - # d 1 +with ( # d 1 x ): pass diff --git a/crates/ruff_python_formatter/tests/snapshots/format@statement__import_from.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@statement__import_from.py.snap index dd2ca4c80f..433e893297 100644 --- a/crates/ruff_python_formatter/tests/snapshots/format@statement__import_from.py.snap +++ b/crates/ruff_python_formatter/tests/snapshots/format@statement__import_from.py.snap @@ -88,7 +88,9 @@ from a import ( # comment bar, ) -from a import bar # comment +from a import ( # comment + bar +) from a import ( bar, diff --git a/crates/ruff_python_formatter/tests/snapshots/format@statement__with.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@statement__with.py.snap index 3140929dc1..09eb21c805 100644 --- a/crates/ruff_python_formatter/tests/snapshots/format@statement__with.py.snap +++ b/crates/ruff_python_formatter/tests/snapshots/format@statement__with.py.snap @@ -122,6 +122,24 @@ with [ dddddddddddddddddddddddddddddddd, ] as example1, aaaaaaaaaaaaaaaaaaaaaaaaaa * bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb * cccccccccccccccccccccccccccc + ddddddddddddddddd as example2, CtxManager222222222222222() as example2: ... + +# Comments on open parentheses +with ( # comment + CtxManager1() as example1, + CtxManager2() as example2, + CtxManager3() as example3, +): + ... + + +with ( # outer comment + ( # inner comment + CtxManager1() + ) as example1, + CtxManager2() as example2, + CtxManager3() as example3, +): + ... ``` ## Output @@ -250,6 +268,25 @@ with [ dddddddddddddddddddddddddddddddd, ] as example1, aaaaaaaaaaaaaaaaaaaaaaaaaa * bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb * cccccccccccccccccccccccccccc + ddddddddddddddddd as example2, CtxManager222222222222222() as example2: ... + +# Comments on open parentheses +with ( # comment + CtxManager1() as example1, + CtxManager2() as example2, + CtxManager3() as example3, +): + ... + + +with ( # outer comment + ( + # inner comment + CtxManager1() + ) as example1, + CtxManager2() as example2, + CtxManager3() as example3, +): + ... ```