diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/import_from.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/import_from.py index 335e91036abf9d..fb9b9916f0a1be 100644 --- a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/import_from.py +++ b/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 diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs index a09a5b1e6d4a15..4fb803b7960811 100644 --- a/crates/ruff_python_formatter/src/comments/placement.rs +++ b/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 diff --git a/crates/ruff_python_formatter/src/statement/stmt_import_from.rs b/crates/ruff_python_formatter/src/statement/stmt_import_from.rs index e6164aa66dff0b..6e18bdcecf6dd3 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_import_from.rs +++ b/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(()) + } } diff --git a/crates/ruff_python_formatter/tests/snapshots/format@statement__import_from.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@statement__import_from.py.snap index 0d8c90572e987a..0e9c7089ae163f 100644 --- a/crates/ruff_python_formatter/tests/snapshots/format@statement__import_from.py.snap +++ b/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 ```