Skip to content

Commit 0c2b88f

Browse files
authored
[flake8-simplify] Further simplify to binary in preview for if-else-block-instead-of-if-exp (SIM108) (#12796)
In most cases we should suggest a ternary operator, but there are three edge cases where a binary operator is more appropriate. Given an if-else block of the form ```python if test: target_var = body_value else: target_var = else_value ``` This PR updates the check for SIM108 to the following: - If `test == body_value` and preview enabled, suggest to replace with `target_var = test or else_value` - If `test == not body_value` and preview enabled, suggest to replace with `target_var = body_value and else_value` - If `not test == body_value` and preview enabled, suggest to replace with `target_var = body_value and else_value` - Otherwise, suggest to replace with `target_var = body_value if test else else_value` Closes #12189.
1 parent cf1a57d commit 0c2b88f

5 files changed

Lines changed: 698 additions & 9 deletions

File tree

crates/ruff_linter/resources/test/fixtures/flake8_simplify/SIM108.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,3 +135,62 @@ def f():
135135
x = 3
136136
else:
137137
x = 5
138+
139+
# SIM108 - should suggest
140+
# z = cond or other_cond
141+
if cond:
142+
z = cond
143+
else:
144+
z = other_cond
145+
146+
# SIM108 - should suggest
147+
# z = cond and other_cond
148+
if not cond:
149+
z = cond
150+
else:
151+
z = other_cond
152+
153+
# SIM108 - should suggest
154+
# z = not cond and other_cond
155+
if cond:
156+
z = not cond
157+
else:
158+
z = other_cond
159+
160+
# SIM108 does not suggest
161+
# a binary option in these cases,
162+
# despite the fact that `bool`
163+
# is a subclass of both `int` and `float`
164+
# so, e.g. `True == 1`.
165+
# (Of course, these specific expressions
166+
# should be simplified for other reasons...)
167+
if True:
168+
z = 1
169+
else:
170+
z = other
171+
172+
if False:
173+
z = 1
174+
else:
175+
z = other
176+
177+
if 1:
178+
z = True
179+
else:
180+
z = other
181+
182+
# SIM108 does not suggest a binary option in this
183+
# case, since we'd be reducing the number of calls
184+
# from Two to one.
185+
if foo():
186+
z = foo()
187+
else:
188+
z = other
189+
190+
# SIM108 does not suggest a binary option in this
191+
# case, since we'd be reducing the number of calls
192+
# from Two to one.
193+
if foo():
194+
z = not foo()
195+
else:
196+
z = other

crates/ruff_linter/src/rules/flake8_simplify/mod.rs

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ mod tests {
99
use test_case::test_case;
1010

1111
use crate::registry::Rule;
12+
use crate::settings::types::PreviewMode;
13+
use crate::settings::LinterSettings;
1214
use crate::test::test_path;
1315
use crate::{assert_messages, settings};
1416

@@ -54,4 +56,22 @@ mod tests {
5456
assert_messages!(snapshot, diagnostics);
5557
Ok(())
5658
}
59+
60+
#[test_case(Rule::IfElseBlockInsteadOfIfExp, Path::new("SIM108.py"))]
61+
fn preview_rules(rule_code: Rule, path: &Path) -> Result<()> {
62+
let snapshot = format!(
63+
"preview__{}_{}",
64+
rule_code.noqa_code(),
65+
path.to_string_lossy()
66+
);
67+
let diagnostics = test_path(
68+
Path::new("flake8_simplify").join(path).as_path(),
69+
&LinterSettings {
70+
preview: PreviewMode::Enabled,
71+
..LinterSettings::for_rule(rule_code)
72+
},
73+
)?;
74+
assert_messages!(snapshot, diagnostics);
75+
Ok(())
76+
}
5777
}

crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_if_exp.rs

Lines changed: 133 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
22
use ruff_macros::{derive_message_formats, violation};
3-
use ruff_python_ast::{self as ast, ElifElseClause, Expr, Stmt};
3+
use ruff_python_ast::comparable::ComparableExpr;
4+
use ruff_python_ast::helpers::contains_effect;
5+
use ruff_python_ast::{self as ast, BoolOp, ElifElseClause, Expr, Stmt};
46
use ruff_python_semantic::analyze::typing::{is_sys_version_block, is_type_checking_block};
57
use ruff_text_size::{Ranged, TextRange};
68

@@ -9,12 +11,15 @@ use crate::fix::edits::fits;
911

1012
/// ## What it does
1113
/// Check for `if`-`else`-blocks that can be replaced with a ternary operator.
14+
/// Moreover, in [preview], check if these ternary expressions can be
15+
/// further simplified to binary expressions.
1216
///
1317
/// ## Why is this bad?
1418
/// `if`-`else`-blocks that assign a value to a variable in both branches can
15-
/// be expressed more concisely by using a ternary operator.
19+
/// be expressed more concisely by using a ternary or binary operator.
1620
///
1721
/// ## Example
22+
///
1823
/// ```python
1924
/// if foo:
2025
/// bar = x
@@ -27,24 +32,51 @@ use crate::fix::edits::fits;
2732
/// bar = x if foo else y
2833
/// ```
2934
///
35+
/// Or, in [preview]:
36+
///
37+
/// ```python
38+
/// if cond:
39+
/// z = cond
40+
/// else:
41+
/// z = other_cond
42+
/// ```
43+
///
44+
/// Use instead:
45+
///
46+
/// ```python
47+
/// z = cond or other_cond
48+
/// ```
49+
///
3050
/// ## References
3151
/// - [Python documentation: Conditional expressions](https://docs.python.org/3/reference/expressions.html#conditional-expressions)
52+
///
53+
/// [preview]: https://docs.astral.sh/ruff/preview/
3254
#[violation]
3355
pub struct IfElseBlockInsteadOfIfExp {
56+
/// The ternary or binary expression to replace the `if`-`else`-block.
3457
contents: String,
58+
/// Whether to use a binary or ternary assignment.
59+
kind: AssignmentKind,
3560
}
3661

3762
impl Violation for IfElseBlockInsteadOfIfExp {
3863
const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
3964

4065
#[derive_message_formats]
4166
fn message(&self) -> String {
42-
let IfElseBlockInsteadOfIfExp { contents } = self;
43-
format!("Use ternary operator `{contents}` instead of `if`-`else`-block")
67+
let IfElseBlockInsteadOfIfExp { contents, kind } = self;
68+
match kind {
69+
AssignmentKind::Ternary => {
70+
format!("Use ternary operator `{contents}` instead of `if`-`else`-block")
71+
}
72+
AssignmentKind::Binary => {
73+
format!("Use binary operator `{contents}` instead of `if`-`else`-block")
74+
}
75+
}
4476
}
4577

4678
fn fix_title(&self) -> Option<String> {
47-
let IfElseBlockInsteadOfIfExp { contents } = self;
79+
let IfElseBlockInsteadOfIfExp { contents, .. } = self;
4880
Some(format!("Replace `if`-`else`-block with `{contents}`"))
4981
}
5082
}
@@ -121,9 +153,59 @@ pub(crate) fn if_else_block_instead_of_if_exp(checker: &mut Checker, stmt_if: &a
121153
return;
122154
}
123155

124-
let target_var = &body_target;
125-
let ternary = ternary(target_var, body_value, test, else_value);
126-
let contents = checker.generator().stmt(&ternary);
156+
// In most cases we should now suggest a ternary operator,
157+
// but there are three edge cases where a binary operator
158+
// is more appropriate.
159+
//
160+
// For the reader's convenience, here is how
161+
// the notation translates to the if-else block:
162+
//
163+
// ```python
164+
// if test:
165+
// target_var = body_value
166+
// else:
167+
// target_var = else_value
168+
// ```
169+
//
170+
// The match statement below implements the following
171+
// logic:
172+
// - If `test == body_value` and preview enabled, replace with `target_var = test or else_value`
173+
// - If `test == not body_value` and preview enabled, replace with `target_var = body_value and else_value`
174+
// - If `not test == body_value` and preview enabled, replace with `target_var = body_value and else_value`
175+
// - Otherwise, replace with `target_var = body_value if test else else_value`
176+
let (contents, assignment_kind) =
177+
match (checker.settings.preview.is_enabled(), test, body_value) {
178+
(true, test_node, body_node)
179+
if ComparableExpr::from(test_node) == ComparableExpr::from(body_node)
180+
&& !contains_effect(test_node, |id| {
181+
checker.semantic().has_builtin_binding(id)
182+
}) =>
183+
{
184+
let target_var = &body_target;
185+
let binary = assignment_binary_or(target_var, body_value, else_value);
186+
(checker.generator().stmt(&binary), AssignmentKind::Binary)
187+
}
188+
(true, test_node, body_node)
189+
if (test_node.as_unary_op_expr().is_some_and(|op_expr| {
190+
op_expr.op.is_not()
191+
&& ComparableExpr::from(&op_expr.operand) == ComparableExpr::from(body_node)
192+
}) || body_node.as_unary_op_expr().is_some_and(|op_expr| {
193+
op_expr.op.is_not()
194+
&& ComparableExpr::from(&op_expr.operand) == ComparableExpr::from(test_node)
195+
})) && !contains_effect(test_node, |id| {
196+
checker.semantic().has_builtin_binding(id)
197+
}) =>
198+
{
199+
let target_var = &body_target;
200+
let binary = assignment_binary_and(target_var, body_value, else_value);
201+
(checker.generator().stmt(&binary), AssignmentKind::Binary)
202+
}
203+
_ => {
204+
let target_var = &body_target;
205+
let ternary = assignment_ternary(target_var, body_value, test, else_value);
206+
(checker.generator().stmt(&ternary), AssignmentKind::Ternary)
207+
}
208+
};
127209

128210
// Don't flag if the resulting expression would exceed the maximum line length.
129211
if !fits(
@@ -139,6 +221,7 @@ pub(crate) fn if_else_block_instead_of_if_exp(checker: &mut Checker, stmt_if: &a
139221
let mut diagnostic = Diagnostic::new(
140222
IfElseBlockInsteadOfIfExp {
141223
contents: contents.clone(),
224+
kind: assignment_kind,
142225
},
143226
stmt_if.range(),
144227
);
@@ -154,7 +237,18 @@ pub(crate) fn if_else_block_instead_of_if_exp(checker: &mut Checker, stmt_if: &a
154237
checker.diagnostics.push(diagnostic);
155238
}
156239

157-
fn ternary(target_var: &Expr, body_value: &Expr, test: &Expr, orelse_value: &Expr) -> Stmt {
240+
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
241+
enum AssignmentKind {
242+
Binary,
243+
Ternary,
244+
}
245+
246+
fn assignment_ternary(
247+
target_var: &Expr,
248+
body_value: &Expr,
249+
test: &Expr,
250+
orelse_value: &Expr,
251+
) -> Stmt {
158252
let node = ast::ExprIf {
159253
test: Box::new(test.clone()),
160254
body: Box::new(body_value.clone()),
@@ -168,3 +262,33 @@ fn ternary(target_var: &Expr, body_value: &Expr, test: &Expr, orelse_value: &Exp
168262
};
169263
node1.into()
170264
}
265+
266+
fn assignment_binary_and(target_var: &Expr, left_value: &Expr, right_value: &Expr) -> Stmt {
267+
let node = ast::ExprBoolOp {
268+
op: BoolOp::And,
269+
values: vec![left_value.clone(), right_value.clone()],
270+
range: TextRange::default(),
271+
};
272+
let node1 = ast::StmtAssign {
273+
targets: vec![target_var.clone()],
274+
value: Box::new(node.into()),
275+
range: TextRange::default(),
276+
};
277+
node1.into()
278+
}
279+
280+
fn assignment_binary_or(target_var: &Expr, left_value: &Expr, right_value: &Expr) -> Stmt {
281+
(ast::StmtAssign {
282+
range: TextRange::default(),
283+
targets: vec![target_var.clone()],
284+
value: Box::new(
285+
(ast::ExprBoolOp {
286+
range: TextRange::default(),
287+
op: BoolOp::Or,
288+
values: vec![left_value.clone(), right_value.clone()],
289+
})
290+
.into(),
291+
),
292+
})
293+
.into()
294+
}

0 commit comments

Comments
 (0)