Format ExprYield/ExprYieldFrom (#5921)

Co-authored-by: Micha Reiser <micha@reiser.io>
This commit is contained in:
qdegraaf 2023-07-21 14:07:51 +02:00 committed by GitHub
parent c3b506fca6
commit 519dbdffaa
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
13 changed files with 373 additions and 185 deletions

View file

@ -1,6 +1,8 @@
use crate::context::PyFormatContext;
use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
use crate::expression::maybe_parenthesize_expression;
use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses, Parenthesize};
use crate::{FormatNodeRule, PyFormatter};
use ruff_formatter::prelude::{space, text};
use ruff_formatter::{write, Buffer, FormatResult};
use ruff_python_ast::node::AnyNodeRef;
use rustpython_parser::ast::ExprYield;
@ -10,16 +12,40 @@ pub struct FormatExprYield;
impl FormatNodeRule<ExprYield> for FormatExprYield {
fn fmt_fields(&self, item: &ExprYield, f: &mut PyFormatter) -> FormatResult<()> {
write!(f, [not_yet_implemented(item)])
let ExprYield { range: _, value } = item;
if let Some(val) = value {
write!(
f,
[
text("yield"),
space(),
maybe_parenthesize_expression(val, item, Parenthesize::IfRequired)
]
)?;
} else {
write!(f, [&text("yield")])?;
}
Ok(())
}
}
impl NeedsParentheses for ExprYield {
fn needs_parentheses(
&self,
_parent: AnyNodeRef,
parent: AnyNodeRef,
_context: &PyFormatContext,
) -> OptionalParentheses {
OptionalParentheses::Multiline
// According to https://docs.python.org/3/reference/grammar.html There are two situations
// where we do not want to always parenthesize a yield expression:
// 1. Right hand side of an assignment, e.g. `x = yield y`
// 2. Yield statement, e.g. `def foo(): yield y`
// We catch situation 1 below. Situation 2 does not need to be handled here as
// FormatStmtExpr, does not add parenthesis
if parent.is_stmt_assign() || parent.is_stmt_ann_assign() || parent.is_stmt_aug_assign() {
OptionalParentheses::Multiline
} else {
OptionalParentheses::Always
}
}
}

View file

@ -1,6 +1,8 @@
use crate::context::PyFormatContext;
use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
use crate::expression::maybe_parenthesize_expression;
use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses, Parenthesize};
use crate::{FormatNodeRule, PyFormatter};
use ruff_formatter::prelude::{space, text};
use ruff_formatter::{write, Buffer, FormatResult};
use ruff_python_ast::node::AnyNodeRef;
use rustpython_parser::ast::ExprYieldFrom;
@ -10,16 +12,37 @@ pub struct FormatExprYieldFrom;
impl FormatNodeRule<ExprYieldFrom> for FormatExprYieldFrom {
fn fmt_fields(&self, item: &ExprYieldFrom, f: &mut PyFormatter) -> FormatResult<()> {
write!(f, [not_yet_implemented(item)])
let ExprYieldFrom { range: _, value } = item;
write!(
f,
[
text("yield from"),
space(),
maybe_parenthesize_expression(value, item, Parenthesize::IfRequired)
]
)?;
Ok(())
}
}
impl NeedsParentheses for ExprYieldFrom {
fn needs_parentheses(
&self,
_parent: AnyNodeRef,
parent: AnyNodeRef,
_context: &PyFormatContext,
) -> OptionalParentheses {
OptionalParentheses::Multiline
// According to https://docs.python.org/3/reference/grammar.html There are two situations
// where we do not want to always parenthesize a yield expression:
// 1. Right hand side of an assignment, e.g. `x = yield y`
// 2. Yield statement, e.g. `def foo(): yield y`
// We catch situation 1 below. Situation 2 does not need to be handled here as
// FormatStmtExpr, does not add parenthesis
if parent.is_stmt_assign() || parent.is_stmt_ann_assign() || parent.is_stmt_aug_assign() {
OptionalParentheses::Multiline
} else {
OptionalParentheses::Always
}
}
}