[ruff] improve handling of intermixed comments inside from-imports

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

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

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:

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<Alias> for FormatAlias {
             name,
             asname,
         } = item;
-        DotDelimitedIdentifier::new(name).fmt(f)?;
+        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, [space(), token("as"), space(), asname.format()])?;
+            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(())
     }
 }