Skip to content

Commit

Permalink
feat: make SIM401 catch ternary operations
Browse files Browse the repository at this point in the history
  • Loading branch information
Flowrey committed Sep 15, 2023
1 parent f9e3ea2 commit 720d476
Show file tree
Hide file tree
Showing 4 changed files with 100 additions and 0 deletions.
3 changes: 3 additions & 0 deletions crates/ruff/resources/test/fixtures/flake8_simplify/SIM401.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,3 +114,6 @@
vars[idx] = a_dict[key]
else:
vars[idx] = "default"

# SIM401 (pattern-3)
var = a_dict[key] if key in a_dict else "default3"
3 changes: 3 additions & 0 deletions crates/ruff/src/checkers/ast/analyze/expression.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1251,6 +1251,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
orelse,
range: _,
}) => {
if checker.enabled(Rule::IfElseBlockInsteadOfDictGet) {
flake8_simplify::rules::if_expr_with_dict(checker, expr, test, body, orelse);
}
if checker.enabled(Rule::IfExprWithTrueFalse) {
flake8_simplify::rules::if_expr_with_true_false(checker, expr, test, body, orelse);
}
Expand Down
79 changes: 79 additions & 0 deletions crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs
Original file line number Diff line number Diff line change
Expand Up @@ -988,3 +988,82 @@ pub(crate) fn use_dict_get_with_default(checker: &mut Checker, stmt_if: &ast::St
}
checker.diagnostics.push(diagnostic);
}

/// SIM401
// var = a_dict[key] if key in a_dict else "default3"
pub(crate) fn if_expr_with_dict(
checker: &mut Checker,
expr: &Expr,
test: &Expr,
body: &Expr,
orelse: &Expr,
) {
let Expr::Compare(ast::ExprCompare {
left: test_key,
ops,
comparators: test_dict,
range: _,
}) = test
else {
return;
};
let [test_dict] = test_dict.as_slice() else {
return;
};

if let [CmpOp::NotIn] = ops[..] {
return;
}

let Expr::Subscript(ast::ExprSubscript {
value: expected_subscript,
slice: expected_slice,
..
}) = body
else {
return;
};

if !compare_expr(&expected_slice.into(), &test_key.into())
|| !compare_expr(&test_dict.into(), &expected_subscript.into())
{
return;
}

let node = orelse.clone();
let node1 = *test_key.clone();
let node2 = ast::ExprAttribute {
value: expected_subscript.clone(),
attr: Identifier::new("get".to_string(), TextRange::default()),
ctx: ExprContext::Load,
range: TextRange::default(),
};
let node3 = ast::ExprCall {
func: Box::new(node2.into()),
arguments: Arguments {
args: vec![node1, node],
keywords: vec![],
range: TextRange::default(),
},
range: TextRange::default(),
};

let contents = checker.generator().expr(&node3.into());

let mut diagnostic = Diagnostic::new(
IfElseBlockInsteadOfDictGet {
contents: contents.clone(),
},
expr.range(),
);
if checker.patch(diagnostic.kind.rule()) {
if !checker.indexer().has_comments(expr, checker.locator()) {
diagnostic.set_fix(Fix::suggested(Edit::range_replacement(
contents,
expr.range(),
)));
}
}

checker.diagnostics.push(diagnostic);
}
Original file line number Diff line number Diff line change
Expand Up @@ -159,4 +159,19 @@ SIM401.py:45:5: SIM401 [*] Use `vars[idx] = a_dict.get(key, "default")` instead
50 47 | ###
51 48 | # Negative cases

SIM401.py:119:7: SIM401 [*] Use `a_dict.get(key, "default3")` instead of an `if` block
|
118 | # SIM401 (pattern-3)
119 | var = a_dict[key] if key in a_dict else "default3"
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ SIM401
|
= help: Replace with `a_dict.get(key, "default3")`

Suggested fix
116 116 | vars[idx] = "default"
117 117 |
118 118 | # SIM401 (pattern-3)
119 |-var = a_dict[key] if key in a_dict else "default3"
119 |+var = a_dict.get(key, "default3")


0 comments on commit 720d476

Please sign in to comment.