From 7842c82a0a81fbc84c5a17ed09604e8c2c9f7e6a Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 1 Aug 2023 14:58:05 -0400 Subject: [PATCH] Preserve end-of-line comments on import-from statements (#6216) ## Summary Ensures that we keep comments at the end-of-line in cases like: ```python from foo import ( # comment bar, ) ``` Closes https://github.com/astral-sh/ruff/issues/6067. --- .../fixtures/ruff/statement/import_from.py | 21 ++++++ .../src/comments/placement.rs | 53 +++++++++++++- .../src/statement/stmt_import_from.rs | 22 +++++- .../format@statement__import_from.py.snap | 73 +++++++++++++++++++ 4 files changed, 163 insertions(+), 6 deletions(-) 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 335e91036a..43376564ab 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,24 @@ 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, +) diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs index 5c2ef1698d..57272b8dda 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; +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 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 e6164aa66d..30b76d57b6 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,12 +35,18 @@ impl FormatNodeRule 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()))) @@ -45,4 +54,13 @@ impl FormatNodeRule 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 0d8c90572e..dd2ca4c80f 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,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, +) ```