-
Notifications
You must be signed in to change notification settings - Fork 1.1k
/
expr_yield.rs
120 lines (106 loc) · 3.63 KB
/
expr_yield.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
use ruff_formatter::write;
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::{Expr, ExprYield, ExprYieldFrom};
use ruff_text_size::{Ranged, TextRange};
use crate::expression::maybe_parenthesize_expression;
use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses, Parenthesize};
use crate::prelude::*;
pub(super) enum AnyExpressionYield<'a> {
Yield(&'a ExprYield),
YieldFrom(&'a ExprYieldFrom),
}
impl<'a> AnyExpressionYield<'a> {
const fn is_yield_from(&self) -> bool {
matches!(self, AnyExpressionYield::YieldFrom(_))
}
fn value(&self) -> Option<&Expr> {
match self {
AnyExpressionYield::Yield(yld) => yld.value.as_deref(),
AnyExpressionYield::YieldFrom(yld) => Some(&yld.value),
}
}
}
impl Ranged for AnyExpressionYield<'_> {
fn range(&self) -> TextRange {
match self {
AnyExpressionYield::Yield(yld) => yld.range(),
AnyExpressionYield::YieldFrom(yld) => yld.range(),
}
}
}
impl NeedsParentheses for AnyExpressionYield<'_> {
fn needs_parentheses(
&self,
parent: AnyNodeRef,
_context: &PyFormatContext,
) -> OptionalParentheses {
// 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
}
}
}
impl<'a> From<&'a ExprYield> for AnyExpressionYield<'a> {
fn from(value: &'a ExprYield) -> Self {
AnyExpressionYield::Yield(value)
}
}
impl<'a> From<&'a ExprYieldFrom> for AnyExpressionYield<'a> {
fn from(value: &'a ExprYieldFrom) -> Self {
AnyExpressionYield::YieldFrom(value)
}
}
impl<'a> From<&AnyExpressionYield<'a>> for AnyNodeRef<'a> {
fn from(value: &AnyExpressionYield<'a>) -> Self {
match value {
AnyExpressionYield::Yield(yld) => AnyNodeRef::ExprYield(yld),
AnyExpressionYield::YieldFrom(yld) => AnyNodeRef::ExprYieldFrom(yld),
}
}
}
impl Format<PyFormatContext<'_>> for AnyExpressionYield<'_> {
fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
let keyword = if self.is_yield_from() {
"yield from"
} else {
"yield"
};
if let Some(val) = self.value() {
write!(
f,
[
token(keyword),
space(),
maybe_parenthesize_expression(val, self, Parenthesize::Optional)
]
)?;
} else {
// ExprYieldFrom always has Some(value) so we should never get a bare `yield from`
write!(f, [&token(keyword)])?;
}
Ok(())
}
}
#[derive(Default)]
pub struct FormatExprYield;
impl FormatNodeRule<ExprYield> for FormatExprYield {
fn fmt_fields(&self, item: &ExprYield, f: &mut PyFormatter) -> FormatResult<()> {
AnyExpressionYield::from(item).fmt(f)
}
}
impl NeedsParentheses for ExprYield {
fn needs_parentheses(
&self,
parent: AnyNodeRef,
context: &PyFormatContext,
) -> OptionalParentheses {
AnyExpressionYield::from(self).needs_parentheses(parent, context)
}
}