diff --git a/crates/ruff/resources/test/fixtures/flake8_simplify/SIM116.py b/crates/ruff/resources/test/fixtures/flake8_simplify/SIM116.py new file mode 100644 index 0000000000..aa50352f78 --- /dev/null +++ b/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 diff --git a/crates/ruff/src/checkers/ast.rs b/crates/ruff/src/checkers/ast.rs index e4841b02f1..6006061878 100644 --- a/crates/ruff/src/checkers/ast.rs +++ b/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, diff --git a/crates/ruff/src/codes.rs b/crates/ruff/src/codes.rs index 186a9a39b5..3597979c61 100644 --- a/crates/ruff/src/codes.rs +++ b/crates/ruff/src/codes.rs @@ -276,6 +276,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option { (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, diff --git a/crates/ruff/src/registry.rs b/crates/ruff/src/registry.rs index 0b64ce4690..2e03a3e7ae 100644 --- a/crates/ruff/src/registry.rs +++ b/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, diff --git a/crates/ruff/src/rules/flake8_simplify/mod.rs b/crates/ruff/src/rules/flake8_simplify/mod.rs index b5f6a5c6a2..d8fc8ca467 100644 --- a/crates/ruff/src/rules/flake8_simplify/mod.rs +++ b/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()); diff --git a/crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs b/crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs index 90f3c7ba63..b62fece84c 100644 --- a/crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs +++ b/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 = 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 } = ¤t.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, diff --git a/crates/ruff/src/rules/flake8_simplify/rules/mod.rs b/crates/ruff/src/rules/flake8_simplify/rules/mod.rs index 0670c06a86..0c4ca995e1 100644 --- a/crates/ruff/src/rules/flake8_simplify/rules/mod.rs +++ b/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, diff --git a/crates/ruff/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM116_SIM116.py.snap b/crates/ruff/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM116_SIM116.py.snap new file mode 100644 index 0000000000..cc3486fdac --- /dev/null +++ b/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: ~ + diff --git a/docs/rules/if-to-dict.md b/docs/rules/if-to-dict.md new file mode 100644 index 0000000000..cf96b4fb5d --- /dev/null +++ b/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") +``` \ No newline at end of file diff --git a/ruff.schema.json b/ruff.schema.json index c926073cf9..c7d899ac3a 100644 --- a/ruff.schema.json +++ b/ruff.schema.json @@ -1981,6 +1981,7 @@ "SIM112", "SIM114", "SIM115", + "SIM116", "SIM117", "SIM118", "SIM2",