Preserve end-of-line comments on import-from statements (#6216)

## Summary

Ensures that we keep comments at the end-of-line in cases like:

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

Closes https://github.com/astral-sh/ruff/issues/6067.
This commit is contained in:
Charlie Marsh 2023-08-01 14:58:05 -04:00 committed by GitHub
parent 9c708d8fc1
commit 7842c82a0a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 163 additions and 6 deletions

View file

@ -1,9 +1,12 @@
use crate::builders::{parenthesize_if_expands, PyFormatterExtensions};
use crate::{AsFormat, FormatNodeRule, PyFormatter};
use ruff_formatter::prelude::{dynamic_text, format_with, space, text};
use ruff_formatter::{write, Buffer, Format, FormatResult};
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::{AsFormat, FormatNodeRule, PyFormatter};
#[derive(Default)]
pub struct FormatStmtImportFrom;
@ -32,12 +35,18 @@ impl FormatNodeRule<StmtImportFrom> for FormatStmtImportFrom {
space(),
]
)?;
if let [name] = names.as_slice() {
// star can't be surrounded by parentheses
if name.name.as_str() == "*" {
return text("*").fmt(f);
}
}
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())))
@ -45,4 +54,13 @@ impl FormatNodeRule<StmtImportFrom> for FormatStmtImportFrom {
});
parenthesize_if_expands(&names).fmt(f)
}
fn fmt_dangling_comments(
&self,
_node: &StmtImportFrom,
_f: &mut PyFormatter,
) -> FormatResult<()> {
// Handled in `fmt_fields`
Ok(())
}
}