Refactor StmtIf: Linter (#5460)

See #5458 for detailed information

I've tried to split this into commits for individual rules as much as
feasible and left a bunch of TODO where i'm not sure about the rules

---------

Co-authored-by: Micha Reiser <[email protected]>

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

@@ -23,7 +23,7 @@
 else:
     b = 2
 
-# OK (false negative)
+# SIM108
 if True:
     pass
 else:

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

@@ -94,3 +94,10 @@ def foo(p):
     errors = 1
 else:
     errors = 1
+
+if a:
+    # Ignore branches with diverging comments because it means we're repeating
+    # the bodies because we have different reasons for each branch
+    x = 1
+elif c:
+    x = 1

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

@@ -38,6 +38,15 @@
 else:
     vars[idx] = "defaultß9💣2ℝ6789ß9💣2ℝ6789ß9💣2ℝ6789ß9💣2ℝ6789ß9💣2ℝ6789"
 
+# SIM401
+if foo():
+    pass
+else:
+    if key in a_dict:
+        vars[idx] = a_dict[key]
+    else:
+        vars[idx] = "default"
+
 ###
 # Negative cases
 ###
@@ -105,12 +114,3 @@
     vars[idx] = a_dict[key]
 else:
     vars[idx] = "default"
-
-# OK (false negative for nested else)
-if foo():
-    pass
-else:
-    if key in a_dict:
-        vars[idx] = a_dict[key]
-    else:
-        vars[idx] = "default"

crates/ruff/resources/test/fixtures/pyflakes/F634.py

@@ -1,6 +1,11 @@
 if (1, 2):
     pass
 
+if (3, 4):
+    pass
+elif foo:
+    pass
+
 for _ in range(5):
     if True:
         pass

crates/ruff/resources/test/fixtures/pyupgrade/UP036_4.py

@@ -7,28 +7,28 @@
 
 if True:
     if foo:
-        pass
+        print()
     elif sys.version_info < (3, 3):
         cmd = [sys.executable, "-m", "test.regrtest"]
 
 if True:
     if foo:
-        pass
+        print()
     elif sys.version_info < (3, 3):
         cmd = [sys.executable, "-m", "test.regrtest"]
     elif foo:
         cmd = [sys.executable, "-m", "test", "-j0"]
 
     if foo:
-        pass
+        print()
     elif sys.version_info < (3, 3):
         cmd = [sys.executable, "-m", "test.regrtest"]
 
     if sys.version_info < (3, 3):
         cmd = [sys.executable, "-m", "test.regrtest"]
 
     if foo:
-        pass
+        print()
     elif sys.version_info < (3, 3):
         cmd = [sys.executable, "-m", "test.regrtest"]
     else:

crates/ruff/resources/test/fixtures/tryceratops/TRY004.py

@@ -230,6 +230,15 @@ def incorrect_multi_conditional(arg1, arg2):
         raise Exception("...")  # should be typeerror
 
 
+def multiple_is_instance_checks(some_arg):
+    if isinstance(some_arg, str):
+        pass
+    elif isinstance(some_arg, int):
+        pass
+    else:
+        raise Exception("...")   # should be typeerror
+
+
 class MyCustomTypeValidation(Exception):
     pass
 

crates/ruff/src/autofix/edits.rs

@@ -190,12 +190,24 @@ fn is_lone_child(child: &Stmt, parent: &Stmt) -> bool {
         }
         Stmt::For(ast::StmtFor { body, orelse, .. })
         | Stmt::AsyncFor(ast::StmtAsyncFor { body, orelse, .. })
-        | Stmt::While(ast::StmtWhile { body, orelse, .. })
-        | Stmt::If(ast::StmtIf { body, orelse, .. }) => {
+        | Stmt::While(ast::StmtWhile { body, orelse, .. }) => {
             if is_only(body, child) || is_only(orelse, child) {
                 return true;
             }
         }
+        Stmt::If(ast::StmtIf {
+            body,
+            elif_else_clauses,
+            ..
+        }) => {
+            if is_only(body, child)
+                || elif_else_clauses
+                    .iter()
+                    .any(|ast::ElifElseClause { body, .. }| is_only(body, child))
+            {
+                return true;
+            }
+        }
         Stmt::Try(ast::StmtTry {
             body,
             handlers,

crates/ruff/src/checkers/ast/mod.rs

@@ -1289,77 +1289,50 @@ where
                     }
                 }
             }
-            Stmt::If(ast::StmtIf {
-                test,
-                body,
-                orelse,
-                range: _,
-            }) => {
+            Stmt::If(
+                stmt_if @ ast::StmtIf {
+                    test,
+                    elif_else_clauses,
+                    ..
+                },
+            ) => {
                 if self.enabled(Rule::IfTuple) {
-                    pyflakes::rules::if_tuple(self, stmt, test);
+                    pyflakes::rules::if_tuple(self, stmt_if);
                 }
                 if self.enabled(Rule::CollapsibleIf) {
                     flake8_simplify::rules::nested_if_statements(
                         self,
-                        stmt,
-                        test,
-                        body,
-                        orelse,
+                        stmt_if,
                         self.semantic.stmt_parent(),
                     );
                 }
                 if self.enabled(Rule::IfWithSameArms) {
-                    flake8_simplify::rules::if_with_same_arms(
-                        self,
-                        stmt,
-                        self.semantic.stmt_parent(),
-                    );
+                    flake8_simplify::rules::if_with_same_arms(self, self.locator, stmt_if);
                 }
                 if self.enabled(Rule::NeedlessBool) {
                     flake8_simplify::rules::needless_bool(self, stmt);
                 }
                 if self.enabled(Rule::IfElseBlockInsteadOfDictLookup) {
-                    flake8_simplify::rules::manual_dict_lookup(
-                        self,
-                        stmt,
-                        test,
-                        body,
-                        orelse,
-                        self.semantic.stmt_parent(),
-                    );
+                    flake8_simplify::rules::manual_dict_lookup(self, stmt_if);
                 }
                 if self.enabled(Rule::IfElseBlockInsteadOfIfExp) {
-                    flake8_simplify::rules::use_ternary_operator(
-                        self,
-                        stmt,
-                        self.semantic.stmt_parent(),
-                    );
+                    flake8_simplify::rules::use_ternary_operator(self, stmt);
                 }
                 if self.enabled(Rule::IfElseBlockInsteadOfDictGet) {
-                    flake8_simplify::rules::use_dict_get_with_default(
-                        self,
-                        stmt,
-                        test,
-                        body,
-                        orelse,
-                        self.semantic.stmt_parent(),
-                    );
+                    flake8_simplify::rules::use_dict_get_with_default(self, stmt_if);
                 }
                 if self.enabled(Rule::TypeCheckWithoutTypeError) {
                     tryceratops::rules::type_check_without_type_error(
                         self,
-                        body,
-                        test,
-                        orelse,
+                        stmt_if,
                         self.semantic.stmt_parent(),
                     );
                 }
                 if self.enabled(Rule::OutdatedVersionBlock) {
-                    pyupgrade::rules::outdated_version_block(self, stmt, test, body, orelse);
+                    pyupgrade::rules::outdated_version_block(self, stmt_if);
                 }
                 if self.enabled(Rule::CollapsibleElseIf) {
-                    if let Some(diagnostic) =
-                        pylint::rules::collapsible_else_if(orelse, self.locator)
+                    if let Some(diagnostic) = pylint::rules::collapsible_else_if(elif_else_clauses)
                     {
                         self.diagnostics.push(diagnostic);
                     }
@@ -2035,7 +2008,7 @@ where
                 stmt_if @ ast::StmtIf {
                     test,
                     body,
-                    orelse,
+                    elif_else_clauses,
                     range: _,
                 },
             ) => {

crates/ruff/src/rules/flake8_bugbear/rules/reuse_of_groupby_generator.rs

@@ -55,8 +55,6 @@ struct GroupNameFinder<'a> {
     /// A flag indicating that the `group_name` variable has been overridden
     /// during the visit.
     overridden: bool,
-    /// A stack of `if` statements.
-    parent_ifs: Vec<&'a Stmt>,
     /// A stack of counters where each counter is itself a list of usage count.
     /// This is used specifically for mutually exclusive statements such as an
     /// `if` or `match`.
@@ -77,7 +75,6 @@ impl<'a> GroupNameFinder<'a> {
             usage_count: 0,
             nested: false,
             overridden: false,
-            parent_ifs: Vec::new(),
             counter_stack: Vec::new(),
             exprs: Vec::new(),
         }
@@ -146,56 +143,28 @@ where
             Stmt::If(ast::StmtIf {
                 test,
                 body,
-                orelse,
+                elif_else_clauses,
                 range: _,
             }) => {
-                // Determine whether we're on an `if` arm (as opposed to an `elif`).
-                let is_if_arm = !self.parent_ifs.iter().any(|parent| {
-                    if let Stmt::If(ast::StmtIf { orelse, .. }) = parent {
-                        orelse.len() == 1 && &orelse[0] == stmt
-                    } else {
-                        false
-                    }
-                });
-
-                if is_if_arm {
-                    // Initialize the vector with the count for current branch.
-                    self.counter_stack.push(vec![0]);
-                } else {
-                    // SAFETY: `unwrap` is safe because we're either in `elif` or
-                    // `else` branch which can come only after an `if` branch.
-                    // When inside an `if` branch, a new vector will be pushed
-                    // onto the stack.
-                    self.counter_stack.last_mut().unwrap().push(0);
-                }
-
-                let has_else = orelse
-                    .first()
-                    .map_or(false, |expr| !matches!(expr, Stmt::If(_)));
+                // base if plus branches
+                let mut if_stack = Vec::with_capacity(1 + elif_else_clauses.len());
+                // Initialize the vector with the count for the if branch.
+                if_stack.push(0);
+                self.counter_stack.push(if_stack);
 
-                self.parent_ifs.push(stmt);
-                if has_else {
-                    // There's no `Stmt::Else`; instead, the `else` contents are directly on
-                    // the `orelse` of the `Stmt::If` node. We want to add a new counter for
-                    // the `orelse` branch, but first, we need to visit the `if` body manually.
-                    self.visit_expr(test);
-                    self.visit_body(body);
+                self.visit_expr(test);
+                self.visit_body(body);
 
-                    // Now, we're in an `else` block.
+                for clause in elif_else_clauses {
                     self.counter_stack.last_mut().unwrap().push(0);
-                    self.visit_body(orelse);
-                } else {
-                    visitor::walk_stmt(self, stmt);
+                    self.visit_elif_else_clause(clause);
                 }
-                self.parent_ifs.pop();
 
-                if is_if_arm {
-                    if let Some(last) = self.counter_stack.pop() {
-                        // This is the max number of group usage from all the
-                        // branches of this `if` statement.
-                        let max_count = last.into_iter().max().unwrap_or(0);
-                        self.increment_usage_count(max_count);
-                    }
+                if let Some(last) = self.counter_stack.pop() {
+                    // This is the max number of group usage from all the
+                    // branches of this `if` statement.
+                    let max_count = last.into_iter().max().unwrap_or(0);
+                    self.increment_usage_count(max_count);
                 }
             }
             Stmt::Match(ast::StmtMatch {