[flake8-simplify]: if-to-dict

README.md

@@ -1309,6 +1309,7 @@ For more, see [flake8-simplify](https://pypi.org/project/flake8-simplify/) on Py
 | SIM112 | use-capital-environment-variables | Use capitalized environment variable `{expected}` instead of `{original}` | 🛠 |
 | SIM114 | [if-with-same-arms](https://beta.ruff.rs/docs/rules/if-with-same-arms/) | Combine `if` branches using logical `or` operator |  |
 | SIM115 | open-file-with-context-handler | Use context handler for opening files |  |
+| SIM116 | [if-to-dict](https://beta.ruff.rs/docs/rules/if-to-dict/) | Use a dictionary instead of consecutive `if` statements | 🛠 |
 | SIM117 | multiple-with-statements | Use a single `with` statement with multiple contexts instead of nested `with` statements | 🛠 |
 | SIM118 | key-in-dict | Use `{key} in {dict}` instead of `{key} in {dict}.keys()` | 🛠 |
 | SIM201 | negate-equal-op | Use `{left} != {right}` instead of `not {left} == {right}` | 🛠 |

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

@@ -0,0 +1,61 @@
+# These SHOULD change
+a = "hello"
+
+if a == "foo":
+    return "bar"
+elif a == "bar":
+    return "baz"
+elif a == "boo":
+    return "ooh"
+else:
+    return 42
+
+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)
+
+if a == 1:
+    return (1, 2, 3)
+elif a == 2:
+    return (4, 5, 6)
+elif a == 3:
+    return (7, 8, 9)
+
+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)
+
+if a == b"one":
+    return 1
+elif a == b"two":
+    return 2
+elif a == b"three":
+    return 3
+
+# These Should NOT change
+if a == "foo":
+    return "bar"
+elif a == "bar":
+    return baz()
+elif a == "boo":
+    return "ooh"
+else:
+    return 42
+
+
+if a == b"one":
+    return 1
+elif b == b"two":
+    return 2
+elif a == b"three":
+    return 3

crates/ruff/src/checkers/ast.rs

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

crates/ruff/src/codes.rs

@@ -275,6 +275,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::IfToDict,
         (Flake8Simplify, "117") => Rule::MultipleWithStatements,
         (Flake8Simplify, "118") => Rule::KeyInDict,
         (Flake8Simplify, "201") => Rule::NegateEqualOp,

crates/ruff/src/registry.rs

@@ -252,6 +252,7 @@ ruff_macros::register_rules!(
     rules::flake8_2020::rules::SysVersionCmpStr10,
     rules::flake8_2020::rules::SysVersionSlice1Referenced,
     // flake8-simplify
+    rules::flake8_simplify::rules::IfToDict,
     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::IfToDict, 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

@@ -2,6 +2,7 @@ use log::error;
 use rustpython_parser::ast::{Cmpop, Constant, Expr, ExprContext, ExprKind, Stmt, StmtKind};
 
 use ruff_macros::{define_violation, derive_message_formats};
+use std::collections::HashMap;
 
 use crate::ast::comparable::{ComparableExpr, ComparableStmt};
 use crate::ast::helpers::{
@@ -13,7 +14,7 @@ use crate::checkers::ast::Checker;
 use crate::fix::Fix;
 use crate::registry::Diagnostic;
 use crate::rules::flake8_simplify::rules::fix_if;
-use crate::violation::{AutofixKind, Availability, Violation};
+use crate::violation::{AlwaysAutofixableViolation, AutofixKind, Availability, Violation};
 
 fn compare_expr(expr1: &ComparableExpr, expr2: &ComparableExpr) -> bool {
     expr1.eq(expr2)
@@ -81,6 +82,40 @@ impl Violation for NeedlessBool {
     }
 }
 
+define_violation!(
+    /// ### What it does
+    /// Checks for three or more consective 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 IfToDict;
+);
+impl AlwaysAutofixableViolation for IfToDict {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        format!("Use a dictionary instead of consecutive `if` statements")
+    }
+
+    fn autofix_title(&self) -> String {
+        format!("Replace if statement with a dictionary")
+    }
+}
+
 define_violation!(
     pub struct UseTernaryOperator {
         pub contents: String,
@@ -569,6 +604,205 @@ pub fn if_with_same_arms(checker: &mut Checker, stmt: &Stmt, parent: Option<&Stm
     }
 }
 
+fn should_proceed_main(test: &Expr, body: &[Stmt], orelse: &[Stmt]) -> bool {
+    if let ExprKind::Compare {
+        left,
+        ops,
+        comparators,
+    } = &test.node
+    {
+        if let ExprKind::Name { .. } = &left.node {
+            if ops.len() == 1 && ops[0] == Cmpop::Eq {
+                if comparators.len() == 1 {
+                    if let ExprKind::Constant { .. } = &comparators[0].node {
+                        if body.len() == 1 {
+                            if let StmtKind::Return { .. } = &body[0].node {
+                                if orelse.len() == 1 {
+                                    if let StmtKind::If { .. } = &orelse[0].node {
+                                        return true;
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+    false
+}
+
+fn should_proceed_child(stmt: &Stmt, var_id: &str) -> bool {
+    if let StmtKind::If { test, body, orelse } = &stmt.node {
+        if let ExprKind::Compare {
+            left,
+            ops,
+            comparators,
+        } = &test.node
+        {
+            if let ExprKind::Name { id, .. } = &left.node {
+                if id == var_id && ops.len() == 1 && ops[0] == Cmpop::Eq {
+                    if comparators.len() == 1 {
+                        if let ExprKind::Constant { .. } = &comparators[0].node {
+                            if body.len() == 1 {
+                                if let StmtKind::Return { .. } = &body[0].node {
+                                    if orelse.len() <= 1 {
+                                        return true;
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+    false
+}
+
+fn constant_to_str(value: &Constant) -> String {
+    match value {
+        Constant::None => "None".to_string(),
+        Constant::Str(s) => format!("{:?}", s),
+        Constant::Int(i) => i.to_string(),
+        Constant::Float(f) => f.to_string(),
+        Constant::Ellipsis => "...".to_string(),
+        Constant::Bytes(b) => {
+            let the_string = String::from_utf8(b.clone()).unwrap();
+            format!("b\"{the_string}\"")
+        }
+        Constant::Complex { real, imag } => format!("{}+{}j", real, imag),
+        Constant::Bool(b) => match b {
+            true => "True".to_string(),
+            false => "False".to_string(),
+        },
+        Constant::Tuple(items) => {
+            let mut result = "(".to_string();
+            for item in items {
+                result.push_str(&constant_to_str(item));
+                result.push_str(", ");
+            }
+            result.push(')');
+            result
+        }
+    }
+}
+
+/// SIM116
+pub fn if_to_dict(checker: &mut Checker, stmt: &Stmt, test: &Expr, body: &[Stmt], orelse: &[Stmt]) {
+    if !should_proceed_main(test, body, orelse) {
+        return;
+    }
+    let variable: String;
+    let mut child: Option<&Stmt> = orelse.get(0);
+    let mut else_value: Option<String> = None;
+    let mut key_value_pairs: HashMap<String, String> = HashMap::new();
+    // This check is useless because we also check above, but it makes Rust happy
+    if let ExprKind::Compare {
+        left, comparators, ..
+    } = &test.node
+    {
+        if let ExprKind::Constant { value, .. } = &comparators[0].node {
+            let key = constant_to_str(value);
+            if let StmtKind::Return { value } = &body[0].node {
+                let final_value = match value {
+                    Some(value) => checker
+                        .locator
+                        .slice_source_code_range(&Range::from_located(value)),
+                    None => return,
+                };
+                key_value_pairs.insert(key, final_value.to_string());
+            }
+        }
+        if let ExprKind::Name { id, .. } = &left.node {
+            variable = id.clone();
+        } else {
+            return;
+        }
+    } else {
+        return;
+    }
+    while child.is_some() {
+        if !should_proceed_child(child.unwrap(), &variable) {
+            return;
+        }
+        if let StmtKind::If { test, body, orelse } = &child.unwrap().node {
+            if let StmtKind::Return { value } = &body[0].node {
+                let clean_value = match value {
+                    Some(item) => item,
+                    None => return,
+                };
+                if let ExprKind::Call { .. } = &clean_value.node {
+                    return;
+                }
+                if let ExprKind::Compare { comparators, .. } = &test.node {
+                    if let ExprKind::Constant {
+                        value: const_val, ..
+                    } = &comparators[0].node
+                    {
+                        let key = constant_to_str(const_val);
+                        let final_value = checker
+                            .locator
+                            .slice_source_code_range(&Range::from_located(clean_value));
+                        key_value_pairs.insert(key, final_value.to_string());
+                    }
+                }
+            } else {
+                return;
+            }
+            // let current = checker.locator.slice_source_code_range(&Range::from_located(child.unwrap()));
+
+            if orelse.len() == 1 {
+                match &orelse[0].node {
+                    StmtKind::If { .. } => {
+                        child = orelse.get(0);
+                    }
+                    StmtKind::Return { value } => {
+                        let final_value = match value {
+                            Some(item) => checker
+                                .locator
+                                .slice_source_code_range(&Range::from_located(item)),
+                            None => "None",
+                        };
+                        else_value = Some(final_value.to_string());
+                        child = None;
+                    }
+                    _ => return,
+                }
+            } else {
+                child = None;
+            }
+        } else {
+            return;
+        }
+    }
+    if key_value_pairs.len() < 3 {
+        return;
+    }
+    let mut new_str = format!(
+        "return {{ {} }}",
+        key_value_pairs
+            .iter()
+            .map(|(k, v)| format!("{}: {}", k, v))
+            .collect::<Vec<String>>()
+            .join(", ")
+    );
+    new_str.push_str(&format!(".get({variable}"));
+    if let Some(else_val) = &else_value {
+        new_str.push_str(&format!(", {}", else_val));
+    }
+    new_str.push(')');
+    let mut diagnostic = Diagnostic::new(IfToDict, Range::from_located(stmt));
+    if checker.patch(diagnostic.kind.rule()) {
+        diagnostic.amend(Fix::replacement(
+            new_str,
+            stmt.location,
+            stmt.end_location.unwrap(),
+        ));
+    }
+    checker.diagnostics.push(diagnostic);
+}
+
 /// SIM401
 pub fn use_dict_get_with_default(
     checker: &mut Checker,

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

@@ -5,8 +5,8 @@ pub use ast_bool_op::{
 pub use ast_expr::{use_capital_environment_variables, UseCapitalEnvironmentVariables};
 pub use ast_for::{convert_for_loop_to_any_all, ConvertLoopToAll, ConvertLoopToAny};
 pub use ast_if::{
-    if_with_same_arms, nested_if_statements, return_bool_condition_directly,
-    use_dict_get_with_default, use_ternary_operator, CollapsibleIf, DictGetWithDefault,
+    if_to_dict, if_with_same_arms, nested_if_statements, return_bool_condition_directly,
+    use_dict_get_with_default, use_ternary_operator, CollapsibleIf, DictGetWithDefault, IfToDict,
     IfWithSameArms, NeedlessBool, UseTernaryOperator,
 };
 pub use ast_ifexp::{