This commit is contained in:
Charlie Marsh 2023-08-16 13:23:09 -04:00
parent 51d69b448c
commit 94ffeba4bc
4 changed files with 153 additions and 18 deletions

View File

@ -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
): ...

View File

@ -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.

View File

@ -33,7 +33,8 @@ impl FormatNodeRule<WithItem> 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,

View File

@ -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,
):
...
```