feat: make SIM401 catch ternary operations

astral-sh · Sep 15, 2023 · da38373 · da38373
1 parent f9e3ea2
commit da38373
Show file tree

Hide file tree

Showing 3 changed files with 85 additions and 0 deletions.
diff --git a/crates/ruff/resources/test/fixtures/flake8_simplify/SIM401.py b/crates/ruff/resources/test/fixtures/flake8_simplify/SIM401.py
@@ -14,6 +14,9 @@
 else:
     var = a_dict[key]
 
+# SIM401 (pattern-3)
+var = a_dict[key] if key in a_dict else "default3"
+
 # OK (default contains effect)
 if key in a_dict:
     var = a_dict[key]

diff --git a/crates/ruff/src/checkers/ast/analyze/expression.rs b/crates/ruff/src/checkers/ast/analyze/expression.rs
@@ -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);
             }

diff --git a/crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs b/crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs
@@ -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);
+}