[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

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