[`flake8-simplify`]: Implement manual-dict-lookup (#2767)

crates/ruff/resources/test/fixtures/flake8_simplify/SIM116.py

@@ -0,0 +1,76 @@
+# Errors
+a = "hello"
+
+# SIM116
+if a == "foo":
+    return "bar"
+elif a == "bar":
+    return "baz"
+elif a == "boo":
+    return "ooh"
+else:
+    return 42
+
+# SIM116
+if a == 1:
+    return (1, 2, 3)
+elif a == 2:
+    return (4, 5, 6)
+elif a == 3:
+    return (7, 8, 9)
+else:
+    return (10, 11, 12)
+
+# SIM116
+if a == 1:
+    return (1, 2, 3)
+elif a == 2:
+    return (4, 5, 6)
+elif a == 3:
+    return (7, 8, 9)
+
+# SIM116
+if a == "hello 'sir'":
+    return (1, 2, 3)
+elif a == 'goodbye "mam"':
+    return (4, 5, 6)
+elif a == """Fairwell 'mister'""":
+    return (7, 8, 9)
+else:
+    return (10, 11, 12)
+
+# SIM116
+if a == b"one":
+    return 1
+elif a == b"two":
+    return 2
+elif a == b"three":
+    return 3
+
+# SIM116
+if a == "hello 'sir'":
+    return ("hello'", 'hi"', 3)
+elif a == 'goodbye "mam"':
+    return (4, 5, 6)
+elif a == """Fairwell 'mister'""":
+    return (7, 8, 9)
+else:
+    return (10, 11, 12)
+
+# OK
+if a == "foo":
+    return "bar"
+elif a == "bar":
+    return baz()
+elif a == "boo":
+    return "ooh"
+else:
+    return 42
+
+# OK
+if a == b"one":
+    return 1
+elif b == b"two":
+    return 2
+elif a == b"three":
+    return 3

crates/ruff/src/checkers/ast.rs

@@ -1565,6 +1565,9 @@ where
                 if self.settings.rules.enabled(&Rule::NeedlessBool) {
                     flake8_simplify::rules::return_bool_condition_directly(self, stmt);
                 }
+                if self.settings.rules.enabled(&Rule::ManualDictLookup) {
+                    flake8_simplify::rules::manual_dict_lookup(self, stmt, test, body, orelse);
+                }
                 if self.settings.rules.enabled(&Rule::UseTernaryOperator) {
                     flake8_simplify::rules::use_ternary_operator(
                         self,

crates/ruff/src/codes.rs

@@ -276,6 +276,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<Rule> {
         (Flake8Simplify, "112") => Rule::UseCapitalEnvironmentVariables,
         (Flake8Simplify, "114") => Rule::IfWithSameArms,
         (Flake8Simplify, "115") => Rule::OpenFileWithContextHandler,
+        (Flake8Simplify, "116") => Rule::ManualDictLookup,
         (Flake8Simplify, "117") => Rule::MultipleWithStatements,
         (Flake8Simplify, "118") => Rule::KeyInDict,
         (Flake8Simplify, "201") => Rule::NegateEqualOp,

crates/ruff/src/registry.rs

@@ -253,6 +253,7 @@ ruff_macros::register_rules!(
     rules::flake8_2020::rules::SysVersionCmpStr10,
     rules::flake8_2020::rules::SysVersionSlice1Referenced,
     // flake8-simplify
+    rules::flake8_simplify::rules::ManualDictLookup,
     rules::flake8_simplify::rules::DuplicateIsinstanceCall,
     rules::flake8_simplify::rules::CollapsibleIf,
     rules::flake8_simplify::rules::NeedlessBool,

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

@@ -37,6 +37,7 @@ mod tests {
     #[test_case(Rule::AndFalse, Path::new("SIM223.py"); "SIM223")]
     #[test_case(Rule::YodaConditions, Path::new("SIM300.py"); "SIM300")]
     #[test_case(Rule::DictGetWithDefault, Path::new("SIM401.py"); "SIM401")]
+    #[test_case(Rule::ManualDictLookup, Path::new("SIM116.py"); "SIM116")]
     #[test_case(Rule::IfWithSameArms, Path::new("SIM114.py"); "SIM114")]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());

crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs

@@ -1,9 +1,10 @@
 use log::error;
+use rustc_hash::FxHashSet;
 use rustpython_parser::ast::{Cmpop, Constant, Expr, ExprContext, ExprKind, Stmt, StmtKind};
 
 use ruff_macros::{define_violation, derive_message_formats};
 
-use crate::ast::comparable::{ComparableExpr, ComparableStmt};
+use crate::ast::comparable::{ComparableConstant, ComparableExpr, ComparableStmt};
 use crate::ast::helpers::{
     contains_call_path, contains_effect, create_expr, create_stmt, first_colon_range, has_comments,
     has_comments_in, unparse_expr, unparse_stmt,
@@ -81,6 +82,36 @@ impl Violation for NeedlessBool {
     }
 }
 
+define_violation!(
+    /// ### What it does
+    /// Checks for three or more consecutive if-statements with direct returns
+    ///
+    /// ### Why is this bad?
+    /// These can be simplified by using a dictionary
+    ///
+    /// ### Example
+    /// ```python
+    /// if x = 1:
+    ///     return "Hello"
+    /// elif x = 2:
+    ///     return "Goodbye"
+    /// else:
+    ///    return "Goodnight"
+    /// ```
+    ///
+    /// Use instead:
+    /// ```python
+    /// return {1: "Hello", 2: "Goodbye"}.get(x, "Goodnight")
+    /// ```
+    pub struct ManualDictLookup;
+);
+impl Violation for ManualDictLookup {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        format!("Use a dictionary instead of consecutive `if` statements")
+    }
+}
+
 define_violation!(
     pub struct UseTernaryOperator {
         pub contents: String,
@@ -569,6 +600,122 @@ pub fn if_with_same_arms(checker: &mut Checker, stmt: &Stmt, parent: Option<&Stm
     }
 }
 
+/// SIM116
+pub fn manual_dict_lookup(
+    checker: &mut Checker,
+    stmt: &Stmt,
+    test: &Expr,
+    body: &[Stmt],
+    orelse: &[Stmt],
+) {
+    // Throughout this rule:
+    // * Each if-statement's test must consist of a constant equality check with the same variable.
+    // * Each if-statement's body must consist of a single `return`.
+    // * Each if-statement's orelse must be either another if-statement or empty.
+    // * The final if-statement's orelse must be empty, or a single `return`.
+    let ExprKind::Compare {
+        left,
+        ops,
+        comparators,
+    } = &test.node else {
+        return;
+    };
+    let ExprKind::Name { id: target, .. } = &left.node else {
+        return;
+    };
+    if body.len() != 1 {
+        return;
+    }
+    if orelse.len() != 1 {
+        return;
+    }
+    if !(ops.len() == 1 && ops[0] == Cmpop::Eq) {
+        return;
+    }
+    if comparators.len() != 1 {
+        return;
+    }
+    let ExprKind::Constant { value: constant, .. } = &comparators[0].node else {
+        return;
+    };
+    let StmtKind::Return { value, .. } = &body[0].node else {
+        return;
+    };
+    if value
+        .as_ref()
+        .map_or(false, |value| contains_effect(checker, value))
+    {
+        return;
+    }
+
+    let mut constants: FxHashSet<ComparableConstant> = FxHashSet::default();
+    constants.insert(constant.into());
+
+    let mut child: Option<&Stmt> = orelse.get(0);
+    while let Some(current) = child.take() {
+        let StmtKind::If { test, body, orelse } = &current.node else {
+            return;
+        };
+        if body.len() != 1 {
+            return;
+        }
+        if orelse.len() > 1 {
+            return;
+        }
+        let ExprKind::Compare {
+            left,
+            ops,
+            comparators,
+        } = &test.node else {
+            return;
+        };
+        let ExprKind::Name { id, .. } = &left.node else {
+            return;
+        };
+        if !(id == target && ops.len() == 1 && ops[0] == Cmpop::Eq) {
+            return;
+        }
+        if comparators.len() != 1 {
+            return;
+        }
+        let ExprKind::Constant { value: constant, .. } = &comparators[0].node else {
+            return;
+        };
+        let StmtKind::Return { value, .. } = &body[0].node else {
+            return;
+        };
+        if value
+            .as_ref()
+            .map_or(false, |value| contains_effect(checker, value))
+        {
+            return;
+        };
+
+        constants.insert(constant.into());
+        if let Some(orelse) = orelse.first() {
+            match &orelse.node {
+                StmtKind::If { .. } => {
+                    child = Some(orelse);
+                }
+                StmtKind::Return { .. } => {
+                    child = None;
+                }
+                _ => return,
+            }
+        } else {
+            child = None;
+        }
+    }
+
+    if constants.len() < 3 {
+        return;
+    }
+
+    checker
+        .diagnostics
+        .push(Diagnostic::new(ManualDictLookup, Range::from_located(stmt)));
+}
+
 /// SIM401
 pub fn use_dict_get_with_default(
     checker: &mut Checker,

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

@@ -4,9 +4,9 @@ pub use ast_bool_op::{
 };
 pub use ast_expr::{use_capital_environment_variables, UseCapitalEnvironmentVariables};
 pub use ast_if::{
-    if_with_same_arms, nested_if_statements, return_bool_condition_directly,
+    if_with_same_arms, manual_dict_lookup, nested_if_statements, return_bool_condition_directly,
     use_dict_get_with_default, use_ternary_operator, CollapsibleIf, DictGetWithDefault,
-    IfWithSameArms, NeedlessBool, UseTernaryOperator,
+    IfWithSameArms, ManualDictLookup, NeedlessBool, UseTernaryOperator,
 };
 pub use ast_ifexp::{
     explicit_false_true_in_ifexpr, explicit_true_false_in_ifexpr, twisted_arms_in_ifexpr,

crates/ruff/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM116_SIM116.py.snap

@@ -0,0 +1,65 @@
+---
+source: crates/ruff/src/rules/flake8_simplify/mod.rs
+expression: diagnostics
+---
+- kind:
+    ManualDictLookup: ~
+  location:
+    row: 5
+    column: 0
+  end_location:
+    row: 12
+    column: 13
+  fix: ~
+  parent: ~
+- kind:
+    ManualDictLookup: ~
+  location:
+    row: 15
+    column: 0
+  end_location:
+    row: 22
+    column: 23
+  fix: ~
+  parent: ~
+- kind:
+    ManualDictLookup: ~
+  location:
+    row: 25
+    column: 0
+  end_location:
+    row: 30
+    column: 20
+  fix: ~
+  parent: ~
+- kind:
+    ManualDictLookup: ~
+  location:
+    row: 33
+    column: 0
+  end_location:
+    row: 40
+    column: 23
+  fix: ~
+  parent: ~
+- kind:
+    ManualDictLookup: ~
+  location:
+    row: 43
+    column: 0
+  end_location:
+    row: 48
+    column: 12
+  fix: ~
+  parent: ~
+- kind:
+    ManualDictLookup: ~
+  location:
+    row: 51
+    column: 0
+  end_location:
+    row: 58
+    column: 23
+  fix: ~
+  parent: ~
+

docs/rules/if-to-dict.md

@@ -0,0 +1,26 @@
+# if-to-dict (SIM116)
+
+Derived from the **flake8-simplify** linter.
+
+Autofix is always available.
+
+### What it does
+Checks for three or more consecutive if-statements with direct returns
+
+### Why is this bad?
+These can be simplified by using a dictionary
+
+### Example
+```python
+if x = 1:
+    return "Hello"
+elif x = 2:
+    return "Goodbye"
+else:
+   return "Goodnight"
+```
+
+Use instead:
+```python
+return {1: "Hello", 2: "Goodbye"}.get(x, "Goodnight")
+```

ruff.schema.json

@@ -1981,6 +1981,7 @@
         "SIM112",
         "SIM114",
         "SIM115",
+        "SIM116",
         "SIM117",
         "SIM118",
         "SIM2",