Format expr generator exp (#5804)

This commit is contained in:
David Szotten 2023-07-19 12:01:58 +01:00 committed by GitHub
parent cda90d071c
commit 5d68ad9008
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
13 changed files with 224 additions and 271 deletions

View file

@ -6,6 +6,7 @@ use ruff_python_ast::node::AnyNodeRef;
use ruff_python_whitespace::{SimpleTokenizer, TokenKind};
use crate::comments::dangling_comments;
use crate::expression::expr_generator_exp::GeneratorExpParentheses;
use crate::expression::parentheses::{
parenthesized, NeedsParentheses, OptionalParentheses, Parentheses,
};
@ -50,13 +51,23 @@ impl FormatNodeRule<ExprCall> for FormatExprCall {
let mut joiner = f.join_comma_separated(item.end());
match args.as_slice() {
[argument] if keywords.is_empty() => {
let parentheses =
if is_single_argument_parenthesized(argument, item.end(), source) {
Parentheses::Always
} else {
Parentheses::Never
};
joiner.entry(argument, &argument.format().with_options(parentheses));
match argument {
Expr::GeneratorExp(generator_exp) => joiner.entry(
generator_exp,
&generator_exp
.format()
.with_options(GeneratorExpParentheses::StripIfOnlyFunctionArg),
),
other => {
let parentheses =
if is_single_argument_parenthesized(argument, item.end(), source) {
Parentheses::Always
} else {
Parentheses::Never
};
joiner.entry(other, &other.format().with_options(parentheses))
}
};
}
arguments => {
joiner

View file

@ -1,21 +1,76 @@
use crate::context::PyFormatContext;
use crate::expression::parentheses::parenthesized;
use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
use crate::{not_yet_implemented_custom_text, FormatNodeRule, PyFormatter};
use ruff_formatter::{write, Buffer, FormatResult};
use crate::prelude::*;
use crate::AsFormat;
use crate::{FormatNodeRule, PyFormatter};
use ruff_formatter::{format_args, write, Buffer, FormatResult, FormatRuleWithOptions};
use ruff_python_ast::node::AnyNodeRef;
use rustpython_parser::ast::ExprGeneratorExp;
#[derive(Eq, PartialEq, Debug, Default)]
pub enum GeneratorExpParentheses {
#[default]
Default,
// skip parens if the generator exp is the only argument to a function, e.g.
// ```python
// all(x for y in z)`
// ```
StripIfOnlyFunctionArg,
}
impl FormatRuleWithOptions<ExprGeneratorExp, PyFormatContext<'_>> for FormatExprGeneratorExp {
type Options = GeneratorExpParentheses;
fn with_options(mut self, options: Self::Options) -> Self {
self.parentheses = options;
self
}
}
#[derive(Default)]
pub struct FormatExprGeneratorExp;
pub struct FormatExprGeneratorExp {
parentheses: GeneratorExpParentheses,
}
impl FormatNodeRule<ExprGeneratorExp> for FormatExprGeneratorExp {
fn fmt_fields(&self, _item: &ExprGeneratorExp, f: &mut PyFormatter) -> FormatResult<()> {
write!(
f,
[not_yet_implemented_custom_text(
"(NOT_YET_IMPLEMENTED_generator_key for NOT_YET_IMPLEMENTED_generator_key in [])"
)]
)
fn fmt_fields(&self, item: &ExprGeneratorExp, f: &mut PyFormatter) -> FormatResult<()> {
let ExprGeneratorExp {
range: _,
elt,
generators,
} = item;
let joined = format_with(|f| {
f.join_with(soft_line_break_or_space())
.entries(generators.iter().formatted())
.finish()
});
if self.parentheses == GeneratorExpParentheses::StripIfOnlyFunctionArg {
write!(
f,
[
group(&elt.format()),
soft_line_break_or_space(),
group(&joined),
]
)
} else {
write!(
f,
[parenthesized(
"(",
&format_args!(
group(&elt.format()),
soft_line_break_or_space(),
group(&joined)
),
")"
)]
)
}
}
}