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(()) } }