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.
This commit is contained in:
Charlie Marsh 2023-08-14 11:11:03 -04:00 committed by GitHub
parent 3711f8ad59
commit c3a9151eb5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 148 additions and 19 deletions

View file

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

View file

@ -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<StmtImportFrom> for FormatStmtImportFrom {
let StmtImportFrom {
module,
names,
range: _,
level,
range: _,
} = item;
let level_str = level
@ -43,16 +43,29 @@ impl FormatNodeRule<StmtImportFrom> 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(

View file

@ -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<StmtWith> 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<StmtWith> 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<StmtWith> for FormatStmtWith {
f,
[
text(":"),
trailing_comments(dangling_comments),
trailing_comments(colon_comments),
block_indent(&item.body.format())
]
)