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 03ab7a0a3b..f1ea3d7ed6 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 @@ -254,3 +254,35 @@ with ( with (foo() as bar, baz() as bop): pass + +# Trailing comments on items, broken by commas and parentheses. +with ( + a + as + # comment + b # leading comment +): ... + + +with ( + a + as + b, # leading comment + c as d +): ... + + +with ( + a + as + b # leading comment + ,c as d +): ... + + +with ( + a as ( + b # leading comment + ) + ,c as d +): ... diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs index 0a50bc0674..4f5b61afc6 100644 --- a/crates/ruff_python_formatter/src/comments/placement.rs +++ b/crates/ruff_python_formatter/src/comments/placement.rs @@ -235,7 +235,7 @@ fn handle_enclosed_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::StmtWith(_) => handle_with_comment(comment, locator), AnyNodeRef::ExprCall(_) => handle_call_comment(comment), AnyNodeRef::ExprConstant(_) => { if let Some(AnyNodeRef::ExprFString(fstring)) = comment.enclosing_parent() { @@ -1565,7 +1565,7 @@ fn handle_import_from_comment<'a>( /// /// For example, given: /// ```python -/// with ( # foo +/// with ( # comment /// CtxManager1() as example1, /// CtxManager2() as example2, /// CtxManager3() as example3, @@ -1577,17 +1577,59 @@ fn handle_import_from_comment<'a>( /// 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, + locator: &Locator, ) -> 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) + if comment.line_position().is_own_line() { + return CommentPlacement::Default(comment); } + + if let Some(preceding) = comment.preceding_node() { + // If a comment trails a `WithItem` with an `as`, consider attaching it to the optional + // variable, rather than to the `WithItem` itself. + // + // For example, given: + // ```python + // with ( + // a + // as + // b # comment + // , c as d + // ): ... + // ``` + // + // The `# comment` should be attached to the optional variable `b`, rather than to the + // `WithItem` itself. + if let AnyNodeRef::WithItem(ast::WithItem { + context_expr: _, + optional_vars: Some(optional_vars), + range: _, + }) = preceding + { + let tokenizer = SimpleTokenizer::new( + locator.contents(), + TextRange::new(optional_vars.end(), comment.start()), + ); + if !tokenizer + .skip_trivia() + .any(|token| token.kind() == SimpleTokenKind::Comma) + { + return CommentPlacement::trailing(optional_vars.as_ref(), comment); + } + } + } else if let Some(following) = comment.following_node() { + // If a comment precedes the first `WithItem`, attach it as dangling: + // ```python + // with ( # comment + // a as b, + // c as d, + // ): ... + // ``` + if comment.start() < following.start() { + return CommentPlacement::dangling(comment.enclosing_node(), comment); + } + } + + CommentPlacement::Default(comment) } /// Handle comments inside comprehensions, e.g. diff --git a/crates/ruff_python_formatter/src/other/with_item.rs b/crates/ruff_python_formatter/src/other/with_item.rs index fd91501c06..bfca663f89 100644 --- a/crates/ruff_python_formatter/src/other/with_item.rs +++ b/crates/ruff_python_formatter/src/other/with_item.rs @@ -33,7 +33,8 @@ impl FormatNodeRule for FormatWithItem { write!(f, [space(), text("as"), space()])?; if trailing_as_comments.is_empty() { - write!(f, [optional_vars.format()])?; + maybe_parenthesize_expression(optional_vars, item, Parenthesize::IfRequired) + .fmt(f)?; } else { write!( f, 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 1c54019b65..f30bf699ac 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 @@ -260,6 +260,38 @@ with ( with (foo() as bar, baz() as bop): pass + +# Trailing comments on items, broken by commas and parentheses. +with ( + a + as + # comment + b # leading comment +): ... + + +with ( + a + as + b, # leading comment + c as d +): ... + + +with ( + a + as + b # leading comment + ,c as d +): ... + + +with ( + a as ( + b # leading comment + ) + ,c as d +): ... ``` ## Output @@ -290,8 +322,8 @@ with ( with ( a as ( # a # as # own line - b - ), # b # comma + b # b + ), # comma c, # c ): # colon ... # body @@ -299,8 +331,8 @@ with ( with a as ( # a # as # own line - bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb -): # b + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb # b +): pass @@ -375,8 +407,8 @@ with ( a # trailing own line comment ) as ( # trailing as same line comment - b -): # trailing b same line comment + b # trailing b same line comment +): ... with ( @@ -533,6 +565,34 @@ with ( with foo() as bar, baz() as bop: pass + +# Trailing comments on items, broken by commas and parentheses. +with a as ( + # comment + b # leading comment +): + ... + + +with ( + a as b, # leading comment + c as d, +): + ... + + +with ( + a as b, # leading comment + c as d, +): + ... + + +with ( + a as b, # leading comment + c as d, +): + ... ```