Generalize comment-after-bracket handling to lists, sets, etc. (#6320)

## Summary

We already support preserving the end-of-line comment in calls and type
parameters, as in:

```python
foo(  # comment
    bar,
)
```

This PR adds the same behavior for lists, sets, comprehensions, etc.,
such that we preserve:

```python
[  # comment
    1,
    2,
    3,
]
```

And related cases.
This commit is contained in:
Charlie Marsh 2023-08-03 21:28:05 -04:00 committed by GitHub
parent d3aa8b4ee0
commit 1d8759d5df
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
20 changed files with 265 additions and 99 deletions

View file

@ -1,10 +1,10 @@
use ruff_python_ast::Ranged;
use ruff_formatter::prelude::tag::Condition;
use ruff_formatter::{format_args, write, Argument, Arguments};
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::Ranged;
use ruff_python_trivia::{first_non_trivia_token, SimpleToken, SimpleTokenKind, SimpleTokenizer};
use crate::comments::{dangling_comments, SourceComment};
use crate::context::{NodeLevel, WithNodeLevel};
use crate::prelude::*;
@ -110,6 +110,26 @@ where
{
FormatParenthesized {
left,
comments: &[],
content: Argument::new(content),
right,
}
}
/// Formats `content` enclosed by the `left` and `right` parentheses, along with any dangling
/// comments that on the parentheses themselves.
pub(crate) fn parenthesized_with_dangling_comments<'content, 'ast, Content>(
left: &'static str,
comments: &'content [SourceComment],
content: &'content Content,
right: &'static str,
) -> FormatParenthesized<'content, 'ast>
where
Content: Format<PyFormatContext<'ast>>,
{
FormatParenthesized {
left,
comments,
content: Argument::new(content),
right,
}
@ -117,6 +137,7 @@ where
pub(crate) struct FormatParenthesized<'content, 'ast> {
left: &'static str,
comments: &'content [SourceComment],
content: Argument<'content, PyFormatContext<'ast>>,
right: &'static str,
}
@ -124,12 +145,22 @@ pub(crate) struct FormatParenthesized<'content, 'ast> {
impl<'ast> Format<PyFormatContext<'ast>> for FormatParenthesized<'_, 'ast> {
fn fmt(&self, f: &mut Formatter<PyFormatContext<'ast>>) -> FormatResult<()> {
let inner = format_with(|f| {
group(&format_args![
text(self.left),
&soft_block_indent(&Arguments::from(&self.content)),
text(self.right)
])
.fmt(f)
if self.comments.is_empty() {
group(&format_args![
text(self.left),
&soft_block_indent(&Arguments::from(&self.content)),
text(self.right)
])
.fmt(f)
} else {
group(&format_args![
text(self.left),
&line_suffix(&dangling_comments(self.comments)),
&group(&soft_block_indent(&Arguments::from(&self.content))),
text(self.right)
])
.fmt(f)
}
});
let current_level = f.context().node_level();