feat: make SIM401 catch ternary operations

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

Hide file tree

Showing 3 changed files with 110 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,107 @@ 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 {
+        // if key
+        left: test_key, // Name(ExprName { range: 250..253, id: "key", ctx: Load })
+        ops,            // [In]
+        comparators: test_dict, // [Name(ExprName { range: 257..263, id: "a_dict", ctx: Load })]
+        range: _,       // 250..263
+    }) = 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, // Name(ExprName { range: 235..241, id: "a_dict", ctx: Load })
+        slice: expected_slice,     // Name(ExprName { range: 242..245, id: "key", ctx: Load })
+        ..
+    }) = body
+    else {
+        return;
+    };
+
+    let Expr::Constant(ast::ExprConstant {
+        value: _default_value, // Str(StringConstant { value: "default3", unicode: false, implicit_concatenated: false })
+        range: _,
+    }) = orelse
+    else {
+        return;
+    };
+
+    if !compare_expr(&expected_slice.into(), &test_key.into())  // "key" == "key"
+        || !compare_expr(&test_dict.into(), &expected_subscript.into())
+    // "a_dict" == "a_dict"
+    {
+        return;
+    }
+
+    let node = orelse.clone(); // "default"
+    let node1 = *test_key.clone(); // "key"
+    let node2 = ast::ExprAttribute {
+        // a_dict.get()
+        value: expected_subscript.clone(),
+        attr: Identifier::new("get".to_string(), TextRange::default()),
+        ctx: ExprContext::Load,
+        range: TextRange::default(),
+    };
+    let node3 = ast::ExprCall {
+        // a_dict.get(key, "default")
+        func: Box::new(node2.into()),
+        arguments: Arguments {
+            args: vec![node1, node],
+            keywords: vec![],
+            range: TextRange::default(),
+        },
+        range: TextRange::default(),
+    };
+
+    // TODO: Find a way to get expexted_var
+    // let node4 = expected_var.clone();
+    let node4 = ast::ExprName {
+        range: TextRange::default(),
+        id: "TODO".into(),
+        ctx: ExprContext::Store,
+    };
+
+    let node5 = ast::StmtAssign {
+        targets: vec![node4.into()],
+        value: Box::new(node3.into()),
+        range: TextRange::default(),
+    };
+    let contents = checker.generator().stmt(&node5.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);
+}