Skip to content

Commit

Permalink
[flake8-simplify] Further simplify to binary in preview for `if-els…
Browse files Browse the repository at this point in the history
…e-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.
  • Loading branch information
dylwil3 committed Aug 10, 2024
1 parent cf1a57d commit 0c2b88f
Show file tree
Hide file tree
Showing 5 changed files with 698 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -135,3 +135,62 @@ def f():
x = 3
else:
x = 5

# SIM108 - should suggest
# z = cond or other_cond
if cond:
z = cond
else:
z = other_cond

# SIM108 - should suggest
# z = cond and other_cond
if not cond:
z = cond
else:
z = other_cond

# SIM108 - should suggest
# z = not cond and other_cond
if cond:
z = not cond
else:
z = other_cond

# SIM108 does not suggest
# a binary option in these cases,
# despite the fact that `bool`
# is a subclass of both `int` and `float`
# so, e.g. `True == 1`.
# (Of course, these specific expressions
# should be simplified for other reasons...)
if True:
z = 1
else:
z = other

if False:
z = 1
else:
z = other

if 1:
z = True
else:
z = other

# SIM108 does not suggest a binary option in this
# case, since we'd be reducing the number of calls
# from Two to one.
if foo():
z = foo()
else:
z = other

# SIM108 does not suggest a binary option in this
# case, since we'd be reducing the number of calls
# from Two to one.
if foo():
z = not foo()
else:
z = other
20 changes: 20 additions & 0 deletions crates/ruff_linter/src/rules/flake8_simplify/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ mod tests {
use test_case::test_case;

use crate::registry::Rule;
use crate::settings::types::PreviewMode;
use crate::settings::LinterSettings;
use crate::test::test_path;
use crate::{assert_messages, settings};

Expand Down Expand Up @@ -54,4 +56,22 @@ mod tests {
assert_messages!(snapshot, diagnostics);
Ok(())
}

#[test_case(Rule::IfElseBlockInsteadOfIfExp, Path::new("SIM108.py"))]
fn preview_rules(rule_code: Rule, path: &Path) -> Result<()> {
let snapshot = format!(
"preview__{}_{}",
rule_code.noqa_code(),
path.to_string_lossy()
);
let diagnostics = test_path(
Path::new("flake8_simplify").join(path).as_path(),
&LinterSettings {
preview: PreviewMode::Enabled,
..LinterSettings::for_rule(rule_code)
},
)?;
assert_messages!(snapshot, diagnostics);
Ok(())
}
}
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::{self as ast, ElifElseClause, Expr, Stmt};
use ruff_python_ast::comparable::ComparableExpr;
use ruff_python_ast::helpers::contains_effect;
use ruff_python_ast::{self as ast, BoolOp, ElifElseClause, Expr, Stmt};
use ruff_python_semantic::analyze::typing::{is_sys_version_block, is_type_checking_block};
use ruff_text_size::{Ranged, TextRange};

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

/// ## What it does
/// Check for `if`-`else`-blocks that can be replaced with a ternary operator.
/// Moreover, in [preview], check if these ternary expressions can be
/// further simplified to binary expressions.
///
/// ## Why is this bad?
/// `if`-`else`-blocks that assign a value to a variable in both branches can
/// be expressed more concisely by using a ternary operator.
/// be expressed more concisely by using a ternary or binary operator.
///
/// ## Example
///
/// ```python
/// if foo:
/// bar = x
Expand All @@ -27,24 +32,51 @@ use crate::fix::edits::fits;
/// bar = x if foo else y
/// ```
///
/// Or, in [preview]:
///
/// ```python
/// if cond:
/// z = cond
/// else:
/// z = other_cond
/// ```
///
/// Use instead:
///
/// ```python
/// z = cond or other_cond
/// ```
///
/// ## References
/// - [Python documentation: Conditional expressions](https://docs.python.org/3/reference/expressions.html#conditional-expressions)
///
/// [preview]: https://docs.astral.sh/ruff/preview/
#[violation]
pub struct IfElseBlockInsteadOfIfExp {
/// The ternary or binary expression to replace the `if`-`else`-block.
contents: String,
/// Whether to use a binary or ternary assignment.
kind: AssignmentKind,
}

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

#[derive_message_formats]
fn message(&self) -> String {
let IfElseBlockInsteadOfIfExp { contents } = self;
format!("Use ternary operator `{contents}` instead of `if`-`else`-block")
let IfElseBlockInsteadOfIfExp { contents, kind } = self;
match kind {
AssignmentKind::Ternary => {
format!("Use ternary operator `{contents}` instead of `if`-`else`-block")
}
AssignmentKind::Binary => {
format!("Use binary operator `{contents}` instead of `if`-`else`-block")
}
}
}

fn fix_title(&self) -> Option<String> {
let IfElseBlockInsteadOfIfExp { contents } = self;
let IfElseBlockInsteadOfIfExp { contents, .. } = self;
Some(format!("Replace `if`-`else`-block with `{contents}`"))
}
}
Expand Down Expand Up @@ -121,9 +153,59 @@ pub(crate) fn if_else_block_instead_of_if_exp(checker: &mut Checker, stmt_if: &a
return;
}

let target_var = &body_target;
let ternary = ternary(target_var, body_value, test, else_value);
let contents = checker.generator().stmt(&ternary);
// In most cases we should now suggest a ternary operator,
// but there are three edge cases where a binary operator
// is more appropriate.
//
// For the reader's convenience, here is how
// the notation translates to the if-else block:
//
// ```python
// if test:
// target_var = body_value
// else:
// target_var = else_value
// ```
//
// The match statement below implements the following
// logic:
// - If `test == body_value` and preview enabled, replace with `target_var = test or else_value`
// - If `test == not body_value` and preview enabled, replace with `target_var = body_value and else_value`
// - If `not test == body_value` and preview enabled, replace with `target_var = body_value and else_value`
// - Otherwise, replace with `target_var = body_value if test else else_value`
let (contents, assignment_kind) =
match (checker.settings.preview.is_enabled(), test, body_value) {
(true, test_node, body_node)
if ComparableExpr::from(test_node) == ComparableExpr::from(body_node)
&& !contains_effect(test_node, |id| {
checker.semantic().has_builtin_binding(id)
}) =>
{
let target_var = &body_target;
let binary = assignment_binary_or(target_var, body_value, else_value);
(checker.generator().stmt(&binary), AssignmentKind::Binary)
}
(true, test_node, body_node)
if (test_node.as_unary_op_expr().is_some_and(|op_expr| {
op_expr.op.is_not()
&& ComparableExpr::from(&op_expr.operand) == ComparableExpr::from(body_node)
}) || body_node.as_unary_op_expr().is_some_and(|op_expr| {
op_expr.op.is_not()
&& ComparableExpr::from(&op_expr.operand) == ComparableExpr::from(test_node)
})) && !contains_effect(test_node, |id| {
checker.semantic().has_builtin_binding(id)
}) =>
{
let target_var = &body_target;
let binary = assignment_binary_and(target_var, body_value, else_value);
(checker.generator().stmt(&binary), AssignmentKind::Binary)
}
_ => {
let target_var = &body_target;
let ternary = assignment_ternary(target_var, body_value, test, else_value);
(checker.generator().stmt(&ternary), AssignmentKind::Ternary)
}
};

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

fn ternary(target_var: &Expr, body_value: &Expr, test: &Expr, orelse_value: &Expr) -> Stmt {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum AssignmentKind {
Binary,
Ternary,
}

fn assignment_ternary(
target_var: &Expr,
body_value: &Expr,
test: &Expr,
orelse_value: &Expr,
) -> Stmt {
let node = ast::ExprIf {
test: Box::new(test.clone()),
body: Box::new(body_value.clone()),
Expand All @@ -168,3 +262,33 @@ fn ternary(target_var: &Expr, body_value: &Expr, test: &Expr, orelse_value: &Exp
};
node1.into()
}

fn assignment_binary_and(target_var: &Expr, left_value: &Expr, right_value: &Expr) -> Stmt {
let node = ast::ExprBoolOp {
op: BoolOp::And,
values: vec![left_value.clone(), right_value.clone()],
range: TextRange::default(),
};
let node1 = ast::StmtAssign {
targets: vec![target_var.clone()],
value: Box::new(node.into()),
range: TextRange::default(),
};
node1.into()
}

fn assignment_binary_or(target_var: &Expr, left_value: &Expr, right_value: &Expr) -> Stmt {
(ast::StmtAssign {
range: TextRange::default(),
targets: vec![target_var.clone()],
value: Box::new(
(ast::ExprBoolOp {
range: TextRange::default(),
op: BoolOp::Or,
values: vec![left_value.clone(), right_value.clone()],
})
.into(),
),
})
.into()
}
Loading

0 comments on commit 0c2b88f

Please sign in to comment.