Revisit comment placement in ast and formatting

Author: amyreese

crates/ruff_python_formatter/resources/test/fixtures/black/cases/import_comments.py

@@ -6,8 +6,14 @@
 from x import a as b, b as c  # comment
 
 # ensure intermixed end- and own-line comments are all preserved
-# and at least kept in their original order, if not their original
-# positions within the import statement
+from x import (  # one
+    # two
+    a  # three
+    # four
+    ,  # five
+    # six
+)  # seven
+
 from x import (  # alpha
     # bravo
     a  # charlie

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() {
@@ -1936,6 +1939,7 @@ fn handle_bracketed_end_of_line_comment<'a>(
 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 +1967,48 @@ 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)
+        }
     }
 }
 
+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

@@ -16,18 +16,45 @@ impl FormatNodeRule<Alias> for FormatAlias {
             name,
             asname,
         } = item;
-        DotDelimitedIdentifier::new(name).fmt(f)?;
+        write!(f, [DotDelimitedIdentifier::new(name)])?;
+
+        let comments = f.context().comments().clone();
+        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 {
-            let comments = f.context().comments().clone();
-            let dangling = comments.dangling(item);
+            write!(f, [token("as"),])?;
+
+            if comments.has_leading(asname) {
+                write!(
+                    f,
+                    [
+                        trailing_comments(comments.leading(asname)),
+                        hard_line_break()
+                    ]
+                )?;
+            } else {
+                write!(f, [space()])?;
+            }
+
+            write!(f, [asname.format()])?;
+        }
+
+        if comments.has_dangling(item) {
             write!(
                 f,
                 [
-                    space(),
-                    token("as"),
-                    space(),
-                    asname.format(),
-                    trailing_comments(dangling),
+                    trailing_comments(comments.dangling(item)),
+                    hard_line_break()
                 ]
             )?;
         }

crates/ruff_python_formatter/tests/snapshots/black_compatibility@cases__import_comments.py.snap

@@ -0,0 +1,128 @@
+---
+source: crates/ruff_python_formatter/tests/fixtures.rs
+input_file: crates/ruff_python_formatter/resources/test/fixtures/black/cases/import_comments.py
+---
+## Input
+
+```python
+# 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
+
+# ensure intermixed end- and own-line 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
+```
+
+## Black Differences
+
+```diff
+--- Black
++++ Ruff
+@@ -6,15 +6,22 @@
+ from x import a as b, b as c  # comment
+ 
+ # ensure intermixed end- and own-line comments are all preserved
+-# and at least kept in their original order, if not their original
+-# positions within the import statement
++from x import (  # one
++    # two
++    a  # three
++    # four
++    ,  # five
++    # six
++)  # seven
++
+ from x import (  # alpha
+     # bravo
+-    a as b,  # charlie
++    a  # charlie
+     # delta
+-    # echo
+-    # foxtrot  # golf
++    as  # echo
++    # foxtrot
++    b  # golf
+     # hotel
+-    # india
++    ,  # india
+     # juliet
+ )  # kilo
+```
+
+## Ruff Output
+
+```python
+# 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
+
+# ensure intermixed end- and own-line 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
+```
+
+## Black Output
+
+```python
+# 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
+
+# ensure intermixed end- and own-line comments are all preserved
+# and at least kept in their original order, if not their original
+# positions within the import statement
+from x import (  # alpha
+    # bravo
+    a as b,  # charlie
+    # delta
+    # echo
+    # foxtrot  # golf
+    # hotel
+    # india
+    # juliet
+)  # kilo
+```