Preserve end-of-line comments on import-from statements

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/import_from.py

@@ -14,3 +14,19 @@
     aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa as sdkjflsdjlahlfd,
 )
 from aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa import *
+
+
+from a import bar  # comment
+
+from a import bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar  # comment
+
+from a import (  # comment
+    bar,
+)
+
+from a import (  # comment
+    bar
+)
+
+from a import bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar
+# comment

crates/ruff_python_formatter/src/comments/placement.rs

@@ -1,17 +1,16 @@
 use std::cmp::Ordering;
 
+use ruff_python_ast::node::{AnyNodeRef, AstNode};
+use ruff_python_ast::whitespace::indentation;
 use ruff_python_ast::{
     self as ast, Arguments, Comprehension, Expr, ExprAttribute, ExprBinOp, ExprIfExp, ExprSlice,
-    ExprStarred, MatchCase, Ranged,
+    ExprStarred, MatchCase, Ranged, StmtImportFrom,
 };
-use ruff_text_size::TextRange;
-
-use ruff_python_ast::node::{AnyNodeRef, AstNode};
-use ruff_python_ast::whitespace::indentation;
 use ruff_python_trivia::{
     indentation_at_offset, PythonWhitespace, SimpleToken, SimpleTokenKind, SimpleTokenizer,
 };
 use ruff_source_file::{Locator, UniversalNewlines};
+use ruff_text_size::TextRange;
 
 use crate::comments::visitor::{CommentPlacement, DecoratedComment};
 use crate::expression::expr_slice::{assign_comment_in_slice, ExprSliceCommentSection};
@@ -43,7 +42,14 @@ pub(super) fn place_comment<'a>(
 
     // Change comment placement depending on the node type. These can be seen as node-specific
     // fixups.
+    if let Some(AnyNodeRef::StmtImportFrom(import_from)) = comment.preceding_node() {
+        return handle_trailing_import_from_comment(comment, import_from, locator);
+    }
+
     match comment.enclosing_node() {
+        AnyNodeRef::StmtImportFrom(import_from) => {
+            handle_enclosed_import_from_comment(comment, import_from, locator)
+        }
         AnyNodeRef::Arguments(arguments) => {
             handle_arguments_separator_comment(comment, arguments, locator)
         }
@@ -1105,6 +1111,71 @@ fn find_only_token_in_range(
     token
 }
 
+/// Attach an enclosed end-of-line comment to a [`StmtImportFrom`].
+///
+/// For example, given:
+/// ```python
+/// from foo import (  # comment
+///    bar,
+/// )
+/// ```
+///
+/// The comment will be attached to the `StmtImportFrom` node as a dangling comment, to ensure
+/// that it remains on the same line as the `StmtImportFrom` itself.
+fn handle_enclosed_import_from_comment<'a>(
+    comment: DecoratedComment<'a>,
+    import_from: &'a StmtImportFrom,
+    locator: &Locator,
+) -> CommentPlacement<'a> {
+    let node = comment.enclosing_node();
+
+    // The comment needs to be on the same line, but before the first member, to triggering on:
+    // ```python
+    // from foo import (bar, baz,  # comment
+    //     qux,
+    // )
+    // ```
+    if comment.line_position().is_end_of_line()
+        && !locator.contains_line_break(TextRange::new(import_from.start(), comment.start()))
+        && import_from
+            .names
+            .first()
+            .map_or(false, |name| comment.end() < name.start())
+    {
+        CommentPlacement::dangling(node, comment)
+    } else {
+        CommentPlacement::Default(comment)
+    }
+}
+
+/// Attach a trailing end-of-line comment to a [`StmtImportFrom`].
+///
+/// For example, given:
+/// ```python
+/// from foo import bar  # comment
+/// ```
+///
+/// The comment will be attached to the `StmtImportFrom` node as a dangling comment, to avoid
+/// moving the comment "off" of the `StmtImportFrom` node if it's expanded to multiple lines.
+fn handle_trailing_import_from_comment<'a>(
+    comment: DecoratedComment<'a>,
+    import_from: &'a StmtImportFrom,
+    locator: &Locator,
+) -> CommentPlacement<'a> {
+    let Some(node) = comment.preceding_node() else {
+        return CommentPlacement::Default(comment);
+    };
+
+    // The comment needs to be on the same line.
+    if comment.line_position().is_end_of_line()
+        && !locator.contains_line_break(TextRange::new(import_from.start(), comment.start()))
+    {
+        CommentPlacement::dangling(node, comment)
+    } else {
+        CommentPlacement::Default(comment)
+    }
+}
+
 // Handle comments inside comprehensions, e.g.
 //
 // ```python

crates/ruff_python_formatter/src/statement/stmt_import_from.rs

@@ -1,9 +1,12 @@
-use crate::builders::{parenthesize_if_expands, PyFormatterExtensions};
-use crate::{AsFormat, FormatNodeRule, PyFormatter};
 use ruff_formatter::prelude::{dynamic_text, format_with, space, text};
 use ruff_formatter::{write, Buffer, Format, FormatResult};
+use ruff_python_ast::node::AstNode;
 use ruff_python_ast::{Ranged, StmtImportFrom};
 
+use crate::builders::{parenthesize_if_expands, PyFormatterExtensions};
+use crate::comments::trailing_comments;
+use crate::{AsFormat, FormatNodeRule, PyFormatter};
+
 #[derive(Default)]
 pub struct FormatStmtImportFrom;
 
@@ -32,6 +35,11 @@ impl FormatNodeRule<StmtImportFrom> for FormatStmtImportFrom {
                 space(),
             ]
         )?;
+
+        let comments = f.context().comments().clone();
+        let dangling_comments = comments.dangling_comments(item.as_any_node_ref());
+        write!(f, [trailing_comments(dangling_comments)])?;
+
         if let [name] = names.as_slice() {
             // star can't be surrounded by parentheses
             if name.name.as_str() == "*" {
@@ -45,4 +53,13 @@ impl FormatNodeRule<StmtImportFrom> for FormatStmtImportFrom {
         });
         parenthesize_if_expands(&names).fmt(f)
     }
+
+    fn fmt_dangling_comments(
+        &self,
+        _node: &StmtImportFrom,
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
+        // Handled in `fmt_fields`
+        Ok(())
+    }
 }

crates/ruff_python_formatter/tests/snapshots/format@statement__import_from.py.snap

@@ -20,6 +20,22 @@ from a import (
     aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa as sdkjflsdjlahlfd,
 )
 from aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa import *
+
+
+from a import bar  # comment
+
+from a import bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar  # comment
+
+from a import (  # comment
+    bar,
+)
+
+from a import (  # comment
+    bar
+)
+
+from a import bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar, bar
+# comment
 ```
 
 ## Output
@@ -40,6 +56,54 @@ from a import (
     aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa as sdkjflsdjlahlfd,
 )
 from aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa import *
+
+
+from a import bar  # comment
+
+from a import (  # comment
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+)
+
+from a import (  # comment
+    bar,
+)
+
+from a import bar  # comment
+
+from a import (
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+    bar,
+)
+# comment
 ```