Add formatter support for call and class definition Arguments (#6274)

## Summary

This PR leverages the `Arguments` AST node introduced in #6259 in the
formatter, which ensures that we correctly handle trailing comments in
calls, like:

```python
f(
  1,
  # comment
)

pass
```

(Previously, this was treated as a leading comment on `pass`.)

This also allows us to unify the argument handling across calls and
class definitions.

## Test Plan

A bunch of new fixture tests, plus improved Black compatibility.
This commit is contained in:
Charlie Marsh 2023-08-02 11:54:22 -04:00 committed by GitHub
parent b095b7204b
commit 4c53bfe896
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
19 changed files with 640 additions and 252 deletions

View file

@ -1,11 +1,8 @@
use ruff_python_ast::{Arguments, Ranged, StmtClassDef};
use ruff_text_size::TextRange;
use ruff_formatter::write;
use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
use ruff_python_ast::{Ranged, StmtClassDef};
use ruff_python_trivia::{lines_after, skip_trailing_trivia};
use crate::comments::trailing_comments;
use crate::expression::parentheses::{parenthesized, Parentheses};
use crate::comments::{leading_comments, trailing_comments};
use crate::prelude::*;
use crate::statement::suite::SuiteKind;
@ -23,40 +20,84 @@ impl FormatNodeRule<StmtClassDef> for FormatStmtClassDef {
decorator_list,
} = item;
f.join_with(hard_line_break())
.entries(decorator_list.iter().formatted())
.finish()?;
let comments = f.context().comments().clone();
if !decorator_list.is_empty() {
hard_line_break().fmt(f)?;
let dangling_comments = comments.dangling_comments(item);
let trailing_definition_comments_start =
dangling_comments.partition_point(|comment| comment.line_position().is_own_line());
let (leading_definition_comments, trailing_definition_comments) =
dangling_comments.split_at(trailing_definition_comments_start);
if let Some(last_decorator) = decorator_list.last() {
f.join_with(hard_line_break())
.entries(decorator_list.iter().formatted())
.finish()?;
if leading_definition_comments.is_empty() {
write!(f, [hard_line_break()])?;
} else {
// Write any leading definition comments (between last decorator and the header)
// while maintaining the right amount of empty lines between the comment
// and the last decorator.
let decorator_end =
skip_trailing_trivia(last_decorator.end(), f.context().source());
let leading_line = if lines_after(decorator_end, f.context().source()) <= 1 {
hard_line_break()
} else {
empty_line()
};
write!(
f,
[leading_line, leading_comments(leading_definition_comments)]
)?;
}
}
write!(f, [text("class"), space(), name.format()])?;
if arguments
.as_ref()
.is_some_and(|Arguments { args, keywords, .. }| {
!(args.is_empty() && keywords.is_empty())
})
{
parenthesized(
"(",
&FormatInheritanceClause {
class_definition: item,
},
")",
)
.fmt(f)?;
if let Some(arguments) = arguments {
// Drop empty parentheses, e.g., in:
// ```python
// class A():
// ...
// ```
//
// However, preserve any dangling end-of-line comments, e.g., in:
// ```python
// class A( # comment
// ):
// ...
//
// If the arguments contain any dangling own-line comments, we retain the parentheses,
// e.g., in:
// ```python
// class A( # comment
// # comment
// ):
// ...
// ```
if arguments.args.is_empty()
&& arguments.keywords.is_empty()
&& comments
.dangling_comments(arguments)
.iter()
.all(|comment| comment.line_position().is_end_of_line())
{
let dangling = comments.dangling_comments(arguments);
write!(f, [trailing_comments(dangling)])?;
} else {
write!(f, [arguments.format()])?;
}
}
let comments = f.context().comments().clone();
let trailing_head_comments = comments.dangling_comments(item);
write!(
f,
[
text(":"),
trailing_comments(trailing_head_comments),
trailing_comments(trailing_definition_comments),
block_indent(&body.format().with_options(SuiteKind::Class))
]
)
@ -71,58 +112,3 @@ impl FormatNodeRule<StmtClassDef> for FormatStmtClassDef {
Ok(())
}
}
struct FormatInheritanceClause<'a> {
class_definition: &'a StmtClassDef,
}
impl Format<PyFormatContext<'_>> for FormatInheritanceClause<'_> {
fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
let StmtClassDef {
arguments:
Some(Arguments {
args: bases,
keywords,
..
}),
name,
body,
..
} = self.class_definition
else {
return Ok(());
};
let source = f.context().source();
let mut joiner = f.join_comma_separated(body.first().unwrap().start());
if let Some((first, rest)) = bases.split_first() {
// Manually handle parentheses for the first expression because the logic in `FormatExpr`
// doesn't know that it should disregard the parentheses of the inheritance clause.
// ```python
// class Test(A) # A is not parenthesized, the parentheses belong to the inheritance clause
// class Test((A)) # A is parenthesized
// ```
// parentheses from the inheritance clause belong to the expression.
let tokenizer = SimpleTokenizer::new(source, TextRange::new(name.end(), first.start()))
.skip_trivia();
let left_paren_count = tokenizer
.take_while(|token| token.kind() == SimpleTokenKind::LParen)
.count();
// Ignore the first parentheses count
let parentheses = if left_paren_count > 1 {
Parentheses::Always
} else {
Parentheses::Never
};
joiner.entry(first, &first.format().with_options(parentheses));
joiner.nodes(rest.iter());
}
joiner.nodes(keywords.iter()).finish()
}
}

View file

@ -44,7 +44,7 @@ impl FormatRule<AnyFunctionDefinition<'_>, PyFormatContext<'_>> for FormatAnyFun
let trailing_definition_comments_start =
dangling_comments.partition_point(|comment| comment.line_position().is_own_line());
let (leading_function_definition_comments, trailing_definition_comments) =
let (leading_definition_comments, trailing_definition_comments) =
dangling_comments.split_at(trailing_definition_comments_start);
if let Some(last_decorator) = item.decorators().last() {
@ -52,10 +52,10 @@ impl FormatRule<AnyFunctionDefinition<'_>, PyFormatContext<'_>> for FormatAnyFun
.entries(item.decorators().iter().formatted())
.finish()?;
if leading_function_definition_comments.is_empty() {
if leading_definition_comments.is_empty() {
write!(f, [hard_line_break()])?;
} else {
// Write any leading function comments (between last decorator and function header)
// Write any leading definition comments (between last decorator and the header)
// while maintaining the right amount of empty lines between the comment
// and the last decorator.
let decorator_end =
@ -69,10 +69,7 @@ impl FormatRule<AnyFunctionDefinition<'_>, PyFormatContext<'_>> for FormatAnyFun
write!(
f,
[
leading_line,
leading_comments(leading_function_definition_comments)
]
[leading_line, leading_comments(leading_definition_comments)]
)?;
}
}