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.
This commit is contained in:
Charlie Marsh 2023-08-01 14:58:05 -04:00 committed by GitHub
parent 9c708d8fc1
commit 7842c82a0a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 163 additions and 6 deletions

View file

@ -14,3 +14,24 @@ from a import (
aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa as sdkjflsdjlahlfd, aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa as sdkjflsdjlahlfd,
) )
from aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa import * 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,
)

View file

@ -1,17 +1,16 @@
use std::cmp::Ordering; use std::cmp::Ordering;
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::whitespace::indentation;
use ruff_python_ast::{ use ruff_python_ast::{
self as ast, Comprehension, Expr, ExprAttribute, ExprBinOp, ExprIfExp, ExprSlice, ExprStarred, self as ast, Comprehension, Expr, ExprAttribute, ExprBinOp, ExprIfExp, ExprSlice, ExprStarred,
MatchCase, Parameters, Ranged, MatchCase, Parameters, Ranged,
}; };
use ruff_text_size::TextRange;
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::whitespace::indentation;
use ruff_python_trivia::{ use ruff_python_trivia::{
indentation_at_offset, PythonWhitespace, SimpleToken, SimpleTokenKind, SimpleTokenizer, indentation_at_offset, PythonWhitespace, SimpleToken, SimpleTokenKind, SimpleTokenizer,
}; };
use ruff_source_file::{Locator, UniversalNewlines}; use ruff_source_file::{Locator, UniversalNewlines};
use ruff_text_size::TextRange;
use crate::comments::visitor::{CommentPlacement, DecoratedComment}; use crate::comments::visitor::{CommentPlacement, DecoratedComment};
use crate::expression::expr_slice::{assign_comment_in_slice, ExprSliceCommentSection}; use crate::expression::expr_slice::{assign_comment_in_slice, ExprSliceCommentSection};
@ -81,6 +80,7 @@ pub(super) fn place_comment<'a>(
AnyNodeRef::StmtFunctionDef(_) | AnyNodeRef::StmtAsyncFunctionDef(_) => { AnyNodeRef::StmtFunctionDef(_) | AnyNodeRef::StmtAsyncFunctionDef(_) => {
handle_leading_function_with_decorators_comment(comment) handle_leading_function_with_decorators_comment(comment)
} }
AnyNodeRef::StmtImportFrom(import_from) => handle_import_from_comment(comment, import_from),
_ => CommentPlacement::Default(comment), _ => CommentPlacement::Default(comment),
} }
} }
@ -1105,6 +1105,51 @@ fn find_only_token_in_range(
token 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. // Handle comments inside comprehensions, e.g.
// //
// ```python // ```python

View file

@ -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::prelude::{dynamic_text, format_with, space, text};
use ruff_formatter::{write, Buffer, Format, FormatResult}; use ruff_formatter::{write, Buffer, Format, FormatResult};
use ruff_python_ast::node::AstNode;
use ruff_python_ast::{Ranged, StmtImportFrom}; 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)] #[derive(Default)]
pub struct FormatStmtImportFrom; pub struct FormatStmtImportFrom;
@ -32,12 +35,18 @@ impl FormatNodeRule<StmtImportFrom> for FormatStmtImportFrom {
space(), space(),
] ]
)?; )?;
if let [name] = names.as_slice() { if let [name] = names.as_slice() {
// star can't be surrounded by parentheses // star can't be surrounded by parentheses
if name.name.as_str() == "*" { if name.name.as_str() == "*" {
return text("*").fmt(f); 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| { let names = format_with(|f| {
f.join_comma_separated(item.end()) f.join_comma_separated(item.end())
.entries(names.iter().map(|name| (name, name.format()))) .entries(names.iter().map(|name| (name, name.format())))
@ -45,4 +54,13 @@ impl FormatNodeRule<StmtImportFrom> for FormatStmtImportFrom {
}); });
parenthesize_if_expands(&names).fmt(f) parenthesize_if_expands(&names).fmt(f)
} }
fn fmt_dangling_comments(
&self,
_node: &StmtImportFrom,
_f: &mut PyFormatter,
) -> FormatResult<()> {
// Handled in `fmt_fields`
Ok(())
}
} }

View file

@ -20,6 +20,27 @@ from a import (
aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa as sdkjflsdjlahlfd, aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa as sdkjflsdjlahlfd,
) )
from aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa import * 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 ## Output
@ -40,6 +61,58 @@ from a import (
aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa as sdkjflsdjlahlfd, aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa as sdkjflsdjlahlfd,
) )
from aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa import * 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,
)
``` ```