Skip to content

Commit e93f378

Browse files
MichaReiserMichaReiser
authored
Refactor whitespace around operator (#4223)
1 parent 2124feb commit e93f378

File tree

5 files changed

+269
-170
lines changed

5 files changed

+269
-170
lines changed

‎crates/ruff/src/rules/pycodestyle/rules/logical_lines/missing_whitespace_after_keyword.rs

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,9 @@
1-
use itertools::Itertools;
2-
use ruff_text_size::TextRange;
3-
41
use crate::checkers::logical_lines::LogicalLinesContext;
52
use crate::rules::pycodestyle::rules::logical_lines::LogicalLine;
63
use ruff_diagnostics::Violation;
74
use ruff_macros::{derive_message_formats, violation};
85
use ruff_python_ast::token_kind::TokenKind;
6+
use ruff_text_size::TextRange;
97

108
#[violation]
119
pub struct MissingWhitespaceAfterKeyword;
@@ -22,7 +20,10 @@ pub(crate) fn missing_whitespace_after_keyword(
2220
line: &LogicalLine,
2321
context: &mut LogicalLinesContext,
2422
) {
25-
for (tok0, tok1) in line.tokens().iter().tuple_windows() {
23+
for window in line.tokens().windows(2) {
24+
let tok0 = &window[0];
25+
let tok1 = &window[1];
26+
2627
let tok0_kind = tok0.kind();
2728
let tok1_kind = tok1.kind();
2829

‎crates/ruff/src/rules/pycodestyle/rules/logical_lines/missing_whitespace_around_operator.rs

Lines changed: 153 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,9 @@
11
use crate::checkers::logical_lines::LogicalLinesContext;
2-
use ruff_diagnostics::Violation;
2+
use crate::rules::pycodestyle::rules::logical_lines::{LogicalLine, LogicalLineToken};
3+
use ruff_diagnostics::{DiagnosticKind, Violation};
34
use ruff_macros::{derive_message_formats, violation};
45
use ruff_python_ast::token_kind::TokenKind;
5-
use ruff_text_size::{TextRange, TextSize};
6-
7-
use crate::rules::pycodestyle::rules::logical_lines::LogicalLine;
6+
use ruff_text_size::TextRange;
87

98
// E225
109
#[violation]
@@ -56,131 +55,179 @@ pub(crate) fn missing_whitespace_around_operator(
5655
line: &LogicalLine,
5756
context: &mut LogicalLinesContext,
5857
) {
59-
#[derive(Copy, Clone, Eq, PartialEq)]
60-
enum NeedsSpace {
61-
Yes,
62-
No,
63-
Unset,
64-
}
65-
66-
let mut needs_space_main = NeedsSpace::No;
67-
let mut needs_space_aux = NeedsSpace::Unset;
68-
let mut prev_end_aux = TextSize::default();
6958
let mut parens = 0u32;
70-
let mut prev_type: TokenKind = TokenKind::EndOfFile;
71-
let mut prev_end = TextSize::default();
59+
let mut prev_token: Option<&LogicalLineToken> = None;
60+
let mut tokens = line.tokens().iter().peekable();
7261

73-
for token in line.tokens() {
62+
while let Some(token) = tokens.next() {
7463
let kind = token.kind();
7564

76-
if kind.is_skip_comment() {
65+
if kind.is_trivia() {
7766
continue;
7867
}
7968

8069
match kind {
8170
TokenKind::Lpar | TokenKind::Lambda => parens += 1,
82-
TokenKind::Rpar => parens -= 1,
71+
TokenKind::Rpar => parens = parens.saturating_sub(1),
8372
_ => {}
8473
};
8574

86-
let needs_space = needs_space_main == NeedsSpace::Yes
87-
|| needs_space_aux != NeedsSpace::Unset
88-
|| prev_end_aux != TextSize::new(0);
89-
if needs_space {
90-
if token.start() > prev_end {
91-
if needs_space_main != NeedsSpace::Yes && needs_space_aux != NeedsSpace::Yes {
75+
let needs_space = if kind == TokenKind::Equal && parens > 0 {
76+
// Allow keyword args or defaults: foo(bar=None).
77+
NeedsSpace::No
78+
} else if kind == TokenKind::Slash {
79+
// Tolerate the "/" operator in function definition
80+
// For more info see PEP570
81+
82+
// `def f(a, /, b):` or `def f(a, b, /):` or `f = lambda a, /:`
83+
// ^ ^ ^
84+
let slash_in_func = matches!(
85+
tokens.peek().map(|t| t.kind()),
86+
Some(TokenKind::Comma | TokenKind::Rpar | TokenKind::Colon)
87+
);
88+
89+
NeedsSpace::from(!slash_in_func)
90+
} else if kind.is_unary() || kind == TokenKind::DoubleStar {
91+
let is_binary = prev_token.map_or(false, |prev_token| {
92+
let prev_kind = prev_token.kind();
93+
94+
// Check if the operator is used as a binary operator.
95+
// Allow unary operators: -123, -x, +1.
96+
// Allow argument unpacking: foo(*args, **kwargs)
97+
matches!(
98+
prev_kind,
99+
TokenKind::Rpar | TokenKind::Rsqb | TokenKind::Rbrace
100+
) || !(prev_kind.is_operator()
101+
|| prev_kind.is_keyword()
102+
|| prev_kind.is_soft_keyword())
103+
});
104+
105+
if is_binary {
106+
if kind == TokenKind::DoubleStar {
107+
// Enforce consistent spacing, but don't enforce whitespaces.
108+
NeedsSpace::Optional
109+
} else {
110+
NeedsSpace::Yes
111+
}
112+
} else {
113+
NeedsSpace::No
114+
}
115+
} else if is_whitespace_needed(kind) {
116+
NeedsSpace::Yes
117+
} else {
118+
NeedsSpace::No
119+
};
120+
121+
if needs_space != NeedsSpace::No {
122+
let has_leading_trivia = prev_token.map_or(true, |prev| {
123+
prev.end() < token.start() || prev.kind().is_trivia()
124+
});
125+
126+
let has_trailing_trivia = tokens.peek().map_or(true, |next| {
127+
token.end() < next.start() || next.kind().is_trivia()
128+
});
129+
130+
match (has_leading_trivia, has_trailing_trivia) {
131+
// Operator with trailing but no leading space, enforce consistent spacing
132+
(false, true) => {
92133
context.push(
93134
MissingWhitespaceAroundOperator,
94-
TextRange::empty(prev_end_aux),
135+
TextRange::empty(token.start()),
95136
);
96137
}
97-
needs_space_main = NeedsSpace::No;
98-
needs_space_aux = NeedsSpace::Unset;
99-
prev_end_aux = TextSize::new(0);
100-
} else if kind == TokenKind::Greater
101-
&& matches!(prev_type, TokenKind::Less | TokenKind::Minus)
102-
{
103-
// Tolerate the "<>" operator, even if running Python 3
104-
// Deal with Python 3's annotated return value "->"
105-
} else if prev_type == TokenKind::Slash
106-
&& matches!(kind, TokenKind::Comma | TokenKind::Rpar | TokenKind::Colon)
107-
|| (prev_type == TokenKind::Rpar && kind == TokenKind::Colon)
108-
{
109-
// Tolerate the "/" operator in function definition
110-
// For more info see PEP570
111-
} else {
112-
if needs_space_main == NeedsSpace::Yes || needs_space_aux == NeedsSpace::Yes {
113-
context.push(MissingWhitespaceAroundOperator, TextRange::empty(prev_end));
114-
} else if prev_type != TokenKind::DoubleStar {
115-
if prev_type == TokenKind::Percent {
116-
context.push(
117-
MissingWhitespaceAroundModuloOperator,
118-
TextRange::empty(prev_end_aux),
119-
);
120-
} else if !prev_type.is_arithmetic() {
121-
context.push(
122-
MissingWhitespaceAroundBitwiseOrShiftOperator,
123-
TextRange::empty(prev_end_aux),
124-
);
125-
} else {
138+
// Operator with leading but no trailing space, enforce consistent spacing.
139+
(true, false) => {
140+
context.push(
141+
MissingWhitespaceAroundOperator,
142+
TextRange::empty(token.end()),
143+
);
144+
}
145+
// Operator with no space, require spaces if it is required by the operator.
146+
(false, false) => {
147+
if needs_space == NeedsSpace::Yes {
126148
context.push(
127-
MissingWhitespaceAroundArithmeticOperator,
128-
TextRange::empty(prev_end_aux),
149+
diagnostic_kind_for_operator(kind),
150+
TextRange::empty(token.start()),
129151
);
130152
}
131153
}
132-
needs_space_main = NeedsSpace::No;
133-
needs_space_aux = NeedsSpace::Unset;
134-
prev_end_aux = TextSize::new(0);
135-
}
136-
} else if (kind.is_operator() || matches!(kind, TokenKind::Name))
137-
&& prev_end != TextSize::default()
138-
{
139-
if kind == TokenKind::Equal && parens > 0 {
140-
// Allow keyword args or defaults: foo(bar=None).
141-
} else if kind.is_whitespace_needed() {
142-
needs_space_main = NeedsSpace::Yes;
143-
needs_space_aux = NeedsSpace::Unset;
144-
prev_end_aux = TextSize::new(0);
145-
} else if kind.is_unary() {
146-
// Check if the operator is used as a binary operator
147-
// Allow unary operators: -123, -x, +1.
148-
// Allow argument unpacking: foo(*args, **kwargs)
149-
if (matches!(
150-
prev_type,
151-
TokenKind::Rpar | TokenKind::Rsqb | TokenKind::Rbrace
152-
)) || (!prev_type.is_operator()
153-
&& !prev_type.is_keyword()
154-
&& !prev_type.is_soft_keyword())
155-
{
156-
needs_space_main = NeedsSpace::Unset;
157-
needs_space_aux = NeedsSpace::Unset;
158-
prev_end_aux = TextSize::new(0);
154+
(true, true) => {
155+
// Operator has leading and trailing space, all good
159156
}
160-
} else if kind.is_whitespace_optional() {
161-
needs_space_main = NeedsSpace::Unset;
162-
needs_space_aux = NeedsSpace::Unset;
163-
prev_end_aux = TextSize::new(0);
164157
}
158+
}
165159

166-
if needs_space_main == NeedsSpace::Unset {
167-
// Surrounding space is optional, but ensure that
168-
// trailing space matches opening space
169-
prev_end_aux = prev_end;
170-
needs_space_aux = if token.start() == prev_end {
171-
NeedsSpace::No
172-
} else {
173-
NeedsSpace::Yes
174-
};
175-
} else if needs_space_main == NeedsSpace::Yes && token.start() == prev_end_aux {
176-
// A needed opening space was not found
177-
context.push(MissingWhitespaceAroundOperator, TextRange::empty(prev_end));
178-
needs_space_main = NeedsSpace::No;
179-
needs_space_aux = NeedsSpace::Unset;
180-
prev_end_aux = TextSize::new(0);
181-
}
160+
prev_token = Some(token);
161+
}
162+
}
163+
164+
#[derive(Copy, Clone, Eq, PartialEq, Debug)]
165+
enum NeedsSpace {
166+
/// Needs a leading and trailing space.
167+
Yes,
168+
169+
/// Doesn't need a leading or trailing space. Or in other words, we don't care how many
170+
/// leading or trailing spaces that token has.
171+
No,
172+
173+
/// Needs consistent leading and trailing spacing. The operator needs spacing if
174+
/// * it has a leading space
175+
/// * it has a trailing space
176+
Optional,
177+
}
178+
179+
impl From<bool> for NeedsSpace {
180+
fn from(value: bool) -> Self {
181+
if value {
182+
NeedsSpace::Yes
183+
} else {
184+
NeedsSpace::No
182185
}
183-
prev_type = kind;
184-
prev_end = token.end();
185186
}
186187
}
188+
189+
fn diagnostic_kind_for_operator(operator: TokenKind) -> DiagnosticKind {
190+
if operator == TokenKind::Percent {
191+
DiagnosticKind::from(MissingWhitespaceAroundModuloOperator)
192+
} else if operator.is_bitwise_or_shift() {
193+
DiagnosticKind::from(MissingWhitespaceAroundBitwiseOrShiftOperator)
194+
} else if operator.is_arithmetic() {
195+
DiagnosticKind::from(MissingWhitespaceAroundArithmeticOperator)
196+
} else {
197+
DiagnosticKind::from(MissingWhitespaceAroundOperator)
198+
}
199+
}
200+
201+
fn is_whitespace_needed(kind: TokenKind) -> bool {
202+
matches!(
203+
kind,
204+
TokenKind::DoubleStarEqual
205+
| TokenKind::StarEqual
206+
| TokenKind::SlashEqual
207+
| TokenKind::DoubleSlashEqual
208+
| TokenKind::PlusEqual
209+
| TokenKind::MinusEqual
210+
| TokenKind::NotEqual
211+
| TokenKind::Less
212+
| TokenKind::Greater
213+
| TokenKind::PercentEqual
214+
| TokenKind::CircumflexEqual
215+
| TokenKind::AmperEqual
216+
| TokenKind::VbarEqual
217+
| TokenKind::EqEqual
218+
| TokenKind::LessEqual
219+
| TokenKind::GreaterEqual
220+
| TokenKind::LeftShiftEqual
221+
| TokenKind::RightShiftEqual
222+
| TokenKind::Equal
223+
| TokenKind::And
224+
| TokenKind::Or
225+
| TokenKind::In
226+
| TokenKind::Is
227+
| TokenKind::Rarrow
228+
| TokenKind::ColonEqual
229+
| TokenKind::Slash
230+
| TokenKind::Percent
231+
) || kind.is_arithmetic()
232+
|| kind.is_bitwise_or_shift()
233+
}

0 commit comments

Comments
 (0)