From 8fb29eafb866c1b25abb3c329ea5df90aae8bd1d Mon Sep 17 00:00:00 2001 From: Amethyst Reese Date: Tue, 7 Oct 2025 08:14:09 -0700 Subject: [PATCH] [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 --- .../fixtures/black/cases/import_comments.py | 47 +++++++++ .../black/cases/import_comments.py.expect | 46 +++++++++ .../src/comments/placement.rs | 96 ++++++++++++++++++- .../ruff_python_formatter/src/other/alias.rs | 86 ++++++++++++++++- 4 files changed, 269 insertions(+), 6 deletions(-) create mode 100644 crates/ruff_python_formatter/resources/test/fixtures/black/cases/import_comments.py create mode 100644 crates/ruff_python_formatter/resources/test/fixtures/black/cases/import_comments.py.expect diff --git a/crates/ruff_python_formatter/resources/test/fixtures/black/cases/import_comments.py b/crates/ruff_python_formatter/resources/test/fixtures/black/cases/import_comments.py new file mode 100644 index 0000000000..d65ea72720 --- /dev/null +++ b/crates/ruff_python_formatter/resources/test/fixtures/black/cases/import_comments.py @@ -0,0 +1,47 @@ +# ensure trailing comments are preserved +import x # comment +from x import a # comment +from x import a, b # comment +from x import a as b # comment +from x import a as b, b as c # comment + +from x import ( + a, # comment +) +from x import ( + a, # comment + b, +) + +# ensure comma is added +from x import ( + a # comment +) + +# follow black style by merging cases without own-line comments +from x import ( + a # alpha + , # beta + b, +) + +# ensure intermixed comments are all preserved +from x import ( # one + # two + a # three + # four + , # five + # six +) # seven + +from x import ( # alpha + # bravo + a # charlie + # delta + as # echo + # foxtrot + b # golf + # hotel + , # india + # juliet +) # kilo diff --git a/crates/ruff_python_formatter/resources/test/fixtures/black/cases/import_comments.py.expect b/crates/ruff_python_formatter/resources/test/fixtures/black/cases/import_comments.py.expect new file mode 100644 index 0000000000..a493a13d91 --- /dev/null +++ b/crates/ruff_python_formatter/resources/test/fixtures/black/cases/import_comments.py.expect @@ -0,0 +1,46 @@ +# ensure trailing comments are preserved +import x # comment +from x import a # comment +from x import a, b # comment +from x import a as b # comment +from x import a as b, b as c # comment + +from x import ( + a, # comment +) +from x import ( + a, # comment + b, +) + +# ensure comma is added +from x import ( + a, # comment +) + +# follow black style by merging cases without own-line comments +from x import ( + a, # alpha # beta + b, +) + +# ensure intermixed comments are all preserved +from x import ( # one + # two + a # three + # four + , # five + # six +) # seven + +from x import ( # alpha + # bravo + a # charlie + # delta + as # echo + # foxtrot + b # golf + # hotel + , # india + # juliet +) # kilo diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs index 0ec2aa6d2f..234752ba53 100644 --- a/crates/ruff_python_formatter/src/comments/placement.rs +++ b/crates/ruff_python_formatter/src/comments/placement.rs @@ -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: diff --git a/crates/ruff_python_formatter/src/other/alias.rs b/crates/ruff_python_formatter/src/other/alias.rs index bce6485e12..9b997565fa 100644 --- a/crates/ruff_python_formatter/src/other/alias.rs +++ b/crates/ruff_python_formatter/src/other/alias.rs @@ -1,6 +1,7 @@ use ruff_formatter::write; use ruff_python_ast::Alias; +use crate::comments::trailing_comments; use crate::other::identifier::DotDelimitedIdentifier; use crate::prelude::*; @@ -15,10 +16,89 @@ impl FormatNodeRule for FormatAlias { name, asname, } = item; - DotDelimitedIdentifier::new(name).fmt(f)?; - if let Some(asname) = asname { - write!(f, [space(), token("as"), space(), asname.format()])?; + write!(f, [DotDelimitedIdentifier::new(name)])?; + + let comments = f.context().comments().clone(); + + // ```python + // from foo import ( + // bar # comment + // as baz, + // ) + // ``` + if comments.has_trailing(name) { + write!( + f, + [ + trailing_comments(comments.trailing(name)), + hard_line_break() + ] + )?; + } else if asname.is_some() { + write!(f, [space()])?; } + + if let Some(asname) = asname { + write!(f, [token("as")])?; + + // ```python + // from foo import ( + // bar as # comment + // baz, + // ) + // ``` + if comments.has_leading(asname) { + write!( + f, + [ + trailing_comments(comments.leading(asname)), + hard_line_break() + ] + )?; + } else { + write!(f, [space()])?; + } + + write!(f, [asname.format()])?; + } + + // Dangling comment between alias and comma on a following line + // ```python + // from foo import ( + // bar # comment + // , + // ) + // ``` + let dangling = comments.dangling(item); + if !dangling.is_empty() { + write!(f, [trailing_comments(comments.dangling(item))])?; + + // Black will move the comma and merge comments if there is no own-line comment between + // the alias and the comma. + // + // Eg: + // ```python + // from foo import ( + // bar # one + // , # two + // ) + // ``` + // + // Will become: + // ```python + // from foo import ( + // bar, # one # two) + // ``` + // + // Only force a hard line break if an own-line dangling comment is present. + if dangling + .iter() + .any(|comment| comment.line_position().is_own_line()) + { + write!(f, [hard_line_break()])?; + } + } + Ok(()) } }