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

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

@@ -14,3 +14,24 @@
     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
+
+from a import \
+    (  # comment
+    bar,
+)

crates/ruff_python_formatter/src/comments/placement.rs

@@ -1,17 +1,16 @@
 use std::cmp::Ordering;
 
+use ruff_python_ast::node::AnyNodeRef;
+use ruff_python_ast::whitespace::indentation;
 use ruff_python_ast::{
     self as ast, Comprehension, Expr, ExprAttribute, ExprBinOp, ExprIfExp, ExprSlice, ExprStarred,
     MatchCase, Parameters, Ranged,
 };
-use ruff_text_size::TextRange;
-
-use ruff_python_ast::node::AnyNodeRef;
-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};
@@ -81,6 +80,7 @@ pub(super) fn place_comment<'a>(
         AnyNodeRef::StmtFunctionDef(_) | AnyNodeRef::StmtAsyncFunctionDef(_) => {
             handle_leading_function_with_decorators_comment(comment)
         }
+        AnyNodeRef::StmtImportFrom(import_from) => handle_import_from_comment(comment, import_from),
         _ => CommentPlacement::Default(comment),
     }
 }
@@ -1105,6 +1105,51 @@ 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_import_from_comment<'a>(
+    comment: DecoratedComment<'a>,
+    import_from: &'a ast::StmtImportFrom,
+) -> 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:
+    // ```python
+    // from foo import (  # comment
+    //     bar,
+    //     baz,
+    //     qux,
+    // )
+    // ```
+    // However, this should _not_ be treated as a dangling comment:
+    // ```python
+    // from foo import (bar,  # comment
+    //     baz,
+    //     qux,
+    // )
+    // ```
+    // Thus, we check whether the comment is an end-of-line comment _between_ the start of the
+    // statement and the first member. If so, the only possible position is immediately following
+    // the open parenthesis.
+    if comment.line_position().is_end_of_line()
+        && import_from.names.first().is_some_and(|first_name| {
+            import_from.start() < comment.start() && comment.start() < first_name.start()
+        })
+    {
+        CommentPlacement::dangling(comment.enclosing_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,17 +35,32 @@ impl FormatNodeRule<StmtImportFrom> for FormatStmtImportFrom {
                 space(),
             ]
         )?;
+
         if let [name] = names.as_slice() {
             // star can't be surrounded by parentheses
             if name.name.as_str() == "*" {
                 return text("*").fmt(f);
             }
         }
+
+        let comments = f.context().comments().clone();
+        let dangling_comments = comments.dangling_comments(item.as_any_node_ref());
+        write!(f, [trailing_comments(dangling_comments)])?;
+
         let names = format_with(|f| {
             f.join_comma_separated(item.end())
                 .entries(names.iter().map(|name| (name, name.format())))
                 .finish()
         });
         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,27 @@ 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
+
+from a import \
+    (  # comment
+    bar,
+)
 ```
 
 ## Output
@@ -40,6 +61,58 @@ 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 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,
+)
 ```