[ruff] improve handling of intermixed comments inside from-imports (#20561)
Resolves a crash when attempting to format code like: ``` from x import (a as # whatever b) ``` Reworks the way comments are associated with nodes when parsing modules, so that all possible comment positions can be retained and reproduced during formatting. Overall follows Black's formatting style for multi-line import statements. Fixes issue #19138
This commit is contained in:
@@ -311,7 +311,10 @@ fn handle_enclosed_comment<'a>(
|
||||
AnyNodeRef::StmtClassDef(class_def) => {
|
||||
handle_leading_class_with_decorators_comment(comment, class_def)
|
||||
}
|
||||
AnyNodeRef::StmtImportFrom(import_from) => handle_import_from_comment(comment, import_from),
|
||||
AnyNodeRef::StmtImportFrom(import_from) => {
|
||||
handle_import_from_comment(comment, import_from, source)
|
||||
}
|
||||
AnyNodeRef::Alias(alias) => handle_alias_comment(comment, alias, source),
|
||||
AnyNodeRef::StmtWith(with_) => handle_with_comment(comment, with_),
|
||||
AnyNodeRef::ExprCall(_) => handle_call_comment(comment),
|
||||
AnyNodeRef::ExprStringLiteral(_) => match comment.enclosing_parent() {
|
||||
@@ -1922,7 +1925,7 @@ fn handle_bracketed_end_of_line_comment<'a>(
|
||||
CommentPlacement::Default(comment)
|
||||
}
|
||||
|
||||
/// Attach an enclosed end-of-line comment to a [`ast::StmtImportFrom`].
|
||||
/// Attach an enclosed comment to a [`ast::StmtImportFrom`].
|
||||
///
|
||||
/// For example, given:
|
||||
/// ```python
|
||||
@@ -1933,9 +1936,37 @@ fn handle_bracketed_end_of_line_comment<'a>(
|
||||
///
|
||||
/// 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.
|
||||
///
|
||||
/// If the comment's preceding node is an alias, and the comment is *before* a comma:
|
||||
/// ```python
|
||||
/// from foo import (
|
||||
/// bar as baz # comment
|
||||
/// ,
|
||||
/// )
|
||||
/// ```
|
||||
///
|
||||
/// The comment will then be attached to the [`ast::Alias`] node as a dangling comment instead,
|
||||
/// to ensure that it retains its position before the comma.
|
||||
///
|
||||
/// Otherwise, if the comment is *after* the comma or before a following alias:
|
||||
/// ```python
|
||||
/// from foo import (
|
||||
/// bar as baz, # comment
|
||||
/// )
|
||||
///
|
||||
/// from foo import (
|
||||
/// bar,
|
||||
/// # comment
|
||||
/// baz,
|
||||
/// )
|
||||
/// ```
|
||||
///
|
||||
/// Then it will retain the default behavior of being attached to the relevant [`ast::Alias`] node
|
||||
/// as either a leading or trailing comment.
|
||||
fn handle_import_from_comment<'a>(
|
||||
comment: DecoratedComment<'a>,
|
||||
import_from: &'a ast::StmtImportFrom,
|
||||
source: &str,
|
||||
) -> CommentPlacement<'a> {
|
||||
// The comment needs to be on the same line, but before the first member. For example, we want
|
||||
// to treat this as a dangling comment:
|
||||
@@ -1963,10 +1994,69 @@ fn handle_import_from_comment<'a>(
|
||||
{
|
||||
CommentPlacement::dangling(comment.enclosing_node(), comment)
|
||||
} else {
|
||||
CommentPlacement::Default(comment)
|
||||
if let Some(SimpleToken {
|
||||
kind: SimpleTokenKind::Comma,
|
||||
..
|
||||
}) = SimpleTokenizer::starts_at(comment.start(), source)
|
||||
.skip_trivia()
|
||||
.next()
|
||||
{
|
||||
// Treat comments before the comma as dangling, after as trailing (default)
|
||||
if let Some(AnyNodeRef::Alias(alias)) = comment.preceding_node() {
|
||||
CommentPlacement::dangling(alias, comment)
|
||||
} else {
|
||||
CommentPlacement::Default(comment)
|
||||
}
|
||||
} else {
|
||||
CommentPlacement::Default(comment)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Attach an enclosed comment to the appropriate [`ast::Identifier`] within an [`ast::Alias`].
|
||||
///
|
||||
/// For example:
|
||||
/// ```python
|
||||
/// from foo import (
|
||||
/// bar # comment
|
||||
/// as baz,
|
||||
/// )
|
||||
/// ```
|
||||
///
|
||||
/// Will attach the comment as a trailing comment on the first name [`ast::Identifier`].
|
||||
///
|
||||
/// Whereas:
|
||||
/// ```python
|
||||
/// from foo import (
|
||||
/// bar as # comment
|
||||
/// baz,
|
||||
/// )
|
||||
/// ```
|
||||
///
|
||||
/// Will attach the comment as a leading comment on the second name [`ast::Identifier`].
|
||||
fn handle_alias_comment<'a>(
|
||||
comment: DecoratedComment<'a>,
|
||||
alias: &'a ruff_python_ast::Alias,
|
||||
source: &str,
|
||||
) -> CommentPlacement<'a> {
|
||||
if let Some(asname) = &alias.asname {
|
||||
if let Some(SimpleToken {
|
||||
kind: SimpleTokenKind::As,
|
||||
range: as_range,
|
||||
}) = SimpleTokenizer::starts_at(alias.name.end(), source)
|
||||
.skip_trivia()
|
||||
.next()
|
||||
{
|
||||
return if comment.start() < as_range.start() {
|
||||
CommentPlacement::trailing(&alias.name, comment)
|
||||
} else {
|
||||
CommentPlacement::leading(asname, comment)
|
||||
};
|
||||
}
|
||||
}
|
||||
CommentPlacement::Default(comment)
|
||||
}
|
||||
|
||||
/// Attach an enclosed end-of-line comment to a [`ast::StmtWith`].
|
||||
///
|
||||
/// For example, given:
|
||||
|
||||
Reference in New Issue
Block a user