Format while Statement (#4810)

This commit is contained in:
Micha Reiser 2023-06-05 10:24:00 +02:00 committed by GitHub
parent d1d06960f0
commit c65f47d7c4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
18 changed files with 555 additions and 145 deletions

View file

@ -6,27 +6,37 @@ use ruff_formatter::{format_args, write, FormatError, SourceCode};
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::prelude::AstNode;
use ruff_text_size::{TextLen, TextRange, TextSize};
use rustpython_parser::ast::Ranged;
/// Formats the leading comments of a node.
pub(crate) fn leading_comments<T>(node: &T) -> FormatLeadingComments
pub(crate) fn leading_node_comments<T>(node: &T) -> FormatLeadingComments
where
T: AstNode,
{
FormatLeadingComments {
node: node.as_any_node_ref(),
}
FormatLeadingComments::Node(node.as_any_node_ref())
}
/// Formats the passed comments as leading comments
pub(crate) const fn leading_comments(comments: &[SourceComment]) -> FormatLeadingComments {
FormatLeadingComments::Comments(comments)
}
#[derive(Copy, Clone, Debug)]
pub(crate) struct FormatLeadingComments<'a> {
node: AnyNodeRef<'a>,
pub(crate) enum FormatLeadingComments<'a> {
Node(AnyNodeRef<'a>),
Comments(&'a [SourceComment]),
}
impl Format<PyFormatContext<'_>> for FormatLeadingComments<'_> {
fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
let comments = f.context().comments().clone();
for comment in comments.leading_comments(self.node) {
let leading_comments = match self {
FormatLeadingComments::Node(node) => comments.leading_comments(*node),
FormatLeadingComments::Comments(comments) => comments,
};
for comment in leading_comments {
let slice = comment.slice();
let lines_after_comment = lines_after(f.context().contents(), slice.end());
@ -42,32 +52,88 @@ impl Format<PyFormatContext<'_>> for FormatLeadingComments<'_> {
}
}
/// Formats the trailing comments of `node`
pub(crate) fn trailing_comments<T>(node: &T) -> FormatTrailingComments
/// Formats the leading `comments` of an alternate branch and ensures that it preserves the right
/// number of empty lines before. The `last_node` is the last node of the preceding body.
///
/// For example, `last_node` is the last statement in the if body when formatting the leading
/// comments of the `else` branch.
pub(crate) fn leading_alternate_branch_comments<'a, T>(
comments: &'a [SourceComment],
last_node: Option<T>,
) -> FormatLeadingAlternateBranchComments<'a>
where
T: AstNode,
T: Into<AnyNodeRef<'a>>,
{
FormatTrailingComments {
node: node.as_any_node_ref(),
FormatLeadingAlternateBranchComments {
comments,
last_node: last_node.map(std::convert::Into::into),
}
}
pub(crate) struct FormatTrailingComments<'a> {
node: AnyNodeRef<'a>,
pub(crate) struct FormatLeadingAlternateBranchComments<'a> {
comments: &'a [SourceComment],
last_node: Option<AnyNodeRef<'a>>,
}
impl Format<PyFormatContext<'_>> for FormatLeadingAlternateBranchComments<'_> {
fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
if let Some(first_leading) = self.comments.first() {
// Leading comments only preserves the lines after the comment but not before.
// Insert the necessary lines.
if lines_before(f.context().contents(), first_leading.slice().start()) > 1 {
write!(f, [empty_line()])?;
}
write!(f, [leading_comments(self.comments)])?;
} else if let Some(last_preceding) = self.last_node {
// The leading comments formatting ensures that it preserves the right amount of lines after
// We need to take care of this ourselves, if there's no leading `else` comment.
if lines_after(f.context().contents(), last_preceding.end()) > 1 {
write!(f, [empty_line()])?;
}
}
Ok(())
}
}
/// Formats the trailing comments of `node`
pub(crate) fn trailing_node_comments<T>(node: &T) -> FormatTrailingComments
where
T: AstNode,
{
FormatTrailingComments::Node(node.as_any_node_ref())
}
/// Formats the passed comments as trailing comments
pub(crate) fn trailing_comments(comments: &[SourceComment]) -> FormatTrailingComments {
FormatTrailingComments::Comments(comments)
}
pub(crate) enum FormatTrailingComments<'a> {
Node(AnyNodeRef<'a>),
Comments(&'a [SourceComment]),
}
impl Format<PyFormatContext<'_>> for FormatTrailingComments<'_> {
fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
let comments = f.context().comments().clone();
let mut has_empty_lines_before = false;
for trailing in comments.trailing_comments(self.node) {
let trailing_comments = match self {
FormatTrailingComments::Node(node) => comments.trailing_comments(*node),
FormatTrailingComments::Comments(comments) => comments,
};
let mut has_trailing_own_line_comment = false;
for trailing in trailing_comments {
let slice = trailing.slice();
let lines_before_comment = lines_before(f.context().contents(), slice.start());
has_empty_lines_before |= lines_before_comment > 0;
has_trailing_own_line_comment |= trailing.position().is_own_line();
if has_trailing_own_line_comment {
let lines_before_comment = lines_before(f.context().contents(), slice.start());
if has_empty_lines_before {
// A trailing comment at the end of a body or list
// ```python
// def test():
@ -105,7 +171,7 @@ impl Format<PyFormatContext<'_>> for FormatTrailingComments<'_> {
}
/// Formats the dangling comments of `node`.
pub(crate) fn dangling_comments<T>(node: &T) -> FormatDanglingComments
pub(crate) fn dangling_node_comments<T>(node: &T) -> FormatDanglingComments
where
T: AstNode,
{
@ -229,7 +295,7 @@ impl Format<PyFormatContext<'_>> for FormatEmptyLines {
_ => write!(f, [empty_line(), empty_line()]),
},
NodeLevel::Statement => match self.lines {
NodeLevel::CompoundStatement => match self.lines {
0 | 1 => write!(f, [hard_line_break()]),
_ => write!(f, [empty_line()]),
},

View file

@ -103,7 +103,10 @@ use crate::comments::debug::{DebugComment, DebugComments};
use crate::comments::map::MultiMap;
use crate::comments::node_key::NodeRefEqualityKey;
use crate::comments::visitor::CommentsVisitor;
pub(crate) use format::{dangling_comments, leading_comments, trailing_comments};
pub(crate) use format::{
dangling_node_comments, leading_alternate_branch_comments, leading_node_comments,
trailing_comments, trailing_node_comments,
};
use ruff_formatter::{SourceCode, SourceCodeSlice};
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::source_code::CommentRanges;
@ -121,8 +124,6 @@ pub(crate) struct SourceComment {
position: CommentTextPosition,
}
#[allow(unused)]
// TODO(micha): Remove after using the new comments infrastructure in the formatter.
impl SourceComment {
/// Returns the location of the comment in the original source code.
/// Allows retrieving the text of the comment.
@ -184,8 +185,6 @@ pub(crate) enum CommentTextPosition {
OwnLine,
}
#[allow(unused)]
// TODO(micha): Remove after using the new comments infrastructure in the formatter.
impl CommentTextPosition {
pub(crate) const fn is_own_line(self) -> bool {
matches!(self, CommentTextPosition::OwnLine)

View file

@ -5,7 +5,7 @@ use crate::trivia::find_first_non_trivia_character_in_range;
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::source_code::Locator;
use ruff_python_ast::whitespace;
use ruff_text_size::{TextRange, TextSize};
use ruff_text_size::{TextLen, TextRange, TextSize};
use rustpython_parser::ast::Ranged;
use std::cmp::Ordering;
@ -20,6 +20,7 @@ pub(super) fn place_comment<'a>(
.or_else(|comment| handle_in_between_bodies_end_of_line_comment(comment, locator))
.or_else(|comment| handle_trailing_body_comment(comment, locator))
.or_else(handle_trailing_end_of_line_body_comment)
.or_else(|comment| handle_trailing_end_of_line_condition_comment(comment, locator))
.or_else(|comment| handle_positional_only_arguments_separator_comment(comment, locator))
.or_else(|comment| {
handle_trailing_binary_expression_left_or_operator_comment(comment, locator)
@ -471,6 +472,91 @@ fn handle_trailing_end_of_line_body_comment(comment: DecoratedComment<'_>) -> Co
}
}
/// Handles end of line comments after the `:` of a condition
///
/// ```python
/// while True: # comment
/// pass
/// ```
///
/// It attaches the comment as dangling comment to the enclosing `while` statement.
fn handle_trailing_end_of_line_condition_comment<'a>(
comment: DecoratedComment<'a>,
locator: &Locator,
) -> CommentPlacement<'a> {
use ruff_python_ast::prelude::*;
// Must be an end of line comment
if comment.text_position().is_own_line() {
return CommentPlacement::Default(comment);
}
// Must be between the condition expression and the first body element
let (Some(preceding), Some(following)) = (comment.preceding_node(), comment.following_node()) else {
return CommentPlacement::Default(comment);
};
let expression_before_colon = match comment.enclosing_node() {
AnyNodeRef::StmtIf(StmtIf { test: expr, .. })
| AnyNodeRef::StmtWhile(StmtWhile { test: expr, .. })
| AnyNodeRef::StmtFor(StmtFor { iter: expr, .. })
| AnyNodeRef::StmtAsyncFor(StmtAsyncFor { iter: expr, .. }) => {
Some(AnyNodeRef::from(expr.as_ref()))
}
AnyNodeRef::StmtWith(StmtWith { items, .. })
| AnyNodeRef::StmtAsyncWith(StmtAsyncWith { items, .. }) => {
items.last().map(AnyNodeRef::from)
}
_ => None,
};
let Some(last_before_colon) = expression_before_colon else {
return CommentPlacement::Default(comment);
};
// If the preceding is the node before the `colon`
// `while true:` The node before the `colon` is the `true` constant.
if preceding.ptr_eq(last_before_colon) {
let mut start = preceding.end();
while let Some((offset, c)) = find_first_non_trivia_character_in_range(
locator.contents(),
TextRange::new(start, following.start()),
) {
match c {
':' => {
if comment.slice().start() > offset {
// Comment comes after the colon
// ```python
// while a: # comment
// ...
// ```
return CommentPlacement::dangling(comment.enclosing_node(), comment);
}
// Comment comes before the colon
// ```python
// while (
// a # comment
// ):
// ...
// ```
break;
}
')' => {
// Skip over any closing parentheses
start = offset + ')'.text_len();
}
_ => {
unreachable!("Only ')' or ':' should follow the condition")
}
}
}
}
CommentPlacement::Default(comment)
}
/// Attaches comments for the positional-only arguments separator `/` as trailing comments to the
/// enclosing [`Arguments`] node.
///

View file

@ -19,19 +19,19 @@ expression: comments.debug(test_case.source_code)
"trailing": [],
},
Node {
kind: ExprCompare,
range: 51..57,
source: `x == y`,
kind: StmtIf,
range: 48..212,
source: `if x == y: # if statement e...ne comment⏎`,
}: {
"leading": [],
"dangling": [],
"trailing": [
"dangling": [
SourceComment {
text: "# if statement end of line comment",
position: EndOfLine,
formatted: false,
},
],
"trailing": [],
},
Node {
kind: StmtIf,