Support mutually exclusive branches for B031 (#3844)

crates/ruff/resources/test/fixtures/flake8_bugbear/B031.py

@@ -78,6 +78,49 @@ def collect_shop_items(shopper, items):
     for shopper in shoppers:
         collect_shop_items(shopper, section_items)  # B031
 
+for _section, section_items in itertools.groupby(items, key=lambda p: p[1]):
+    # Mutually exclusive branches shouldn't trigger the warning
+    if _section == "greens":
+        collect_shop_items(shopper, section_items)
+        if _section == "greens":
+            collect_shop_items(shopper, section_items)  # B031
+        elif _section == "frozen items":
+            collect_shop_items(shopper, section_items)  # B031
+        else:
+            collect_shop_items(shopper, section_items)  # B031
+        collect_shop_items(shopper, section_items)  # B031
+    elif _section == "frozen items":
+        # Mix `match` and `if` statements
+        match shopper:
+            case "Jane":
+                collect_shop_items(shopper, section_items)
+                if _section == "fourth":
+                    collect_shop_items(shopper, section_items)  # B031
+            case _:
+                collect_shop_items(shopper, section_items)
+    else:
+        collect_shop_items(shopper, section_items)
+    # Now, it should detect
+    collect_shop_items(shopper, section_items)  # B031
+
+for _section, section_items in itertools.groupby(items, key=lambda p: p[1]):
+    # Mutually exclusive branches shouldn't trigger the warning
+    match _section:
+        case "greens":
+            collect_shop_items(shopper, section_items)
+            match shopper:
+                case "Jane":
+                    collect_shop_items(shopper, section_items)  # B031
+                case _:
+                    collect_shop_items(shopper, section_items)  # B031
+        case "frozen items":
+            collect_shop_items(shopper, section_items)
+            collect_shop_items(shopper, section_items)  # B031
+        case _:
+            collect_shop_items(shopper, section_items)
+    # Now, it should detect
+    collect_shop_items(shopper, section_items)  # B031
+
 for group in groupby(items, key=lambda p: p[1]):
     # This is bad, but not detected currently
     collect_shop_items("Jane", group[1])

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

@@ -49,13 +49,24 @@ struct GroupNameFinder<'a> {
     /// Variable name for the group.
     group_name: &'a str,
     /// Number of times the `group_name` variable was seen during the visit.
-    usage_count: u8,
+    usage_count: u32,
     /// A flag indicating that the visitor is inside a nested `for` or `while`
     /// loop or inside a `dict`, `list` or `set` comprehension.
     nested: bool,
     /// 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`.
+    ///
+    /// The stack element represents an entire `if` or `match` statement while
+    /// the number inside the element represents the usage count for one of
+    /// the branches of the statement. The order of the count corresponds the
+    /// branch order.
+    counter_stack: Vec<Vec<u32>>,
     /// A list of reused expressions.
     exprs: Vec<&'a Expr>,
 }
@@ -67,6 +78,8 @@ impl<'a> GroupNameFinder<'a> {
             usage_count: 0,
             nested: false,
             overridden: false,
+            parent_ifs: Vec::new(),
+            counter_stack: Vec::new(),
             exprs: Vec::new(),
         }
     }
@@ -119,6 +132,79 @@ where
                 visitor::walk_body(self, body);
                 self.nested = false;
             }
+            StmtKind::If { test, body, orelse } => {
+                // 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 StmtKind::If { orelse, .. } = &parent.node {
+                        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 = !is_if_arm
+                    && orelse
+                        .first()
+                        .map_or(false, |expr| !matches!(expr.node, StmtKind::If { .. }));
+
+                self.parent_ifs.push(stmt);
+                if has_else {
+                    // There's no `StmtKind::Else`; instead, the `else` contents are directly on
+                    // the `orelse` of the `StmtKind::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);
+
+                    // Now, we're in an `else` block.
+                    self.counter_stack.last_mut().unwrap().push(0);
+                    self.visit_body(orelse);
+                } else {
+                    visitor::walk_stmt(self, stmt);
+                }
+                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);
+                        if let Some(current_last) = self.counter_stack.last_mut() {
+                            *current_last.last_mut().unwrap() += max_count;
+                        } else {
+                            self.usage_count += max_count;
+                        }
+                    }
+                }
+            }
+            StmtKind::Match { subject, cases } => {
+                self.counter_stack.push(Vec::with_capacity(cases.len()));
+                self.visit_expr(subject);
+                for match_case in cases {
+                    self.counter_stack.last_mut().unwrap().push(0);
+                    self.visit_match_case(match_case);
+                }
+                if let Some(last) = self.counter_stack.pop() {
+                    // This is the max number of group usage from all the
+                    // branches of this `match` statement.
+                    let max_count = last.into_iter().max().unwrap_or(0);
+                    if let Some(current_last) = self.counter_stack.last_mut() {
+                        *current_last.last_mut().unwrap() += max_count;
+                    } else {
+                        self.usage_count += max_count;
+                    }
+                }
+            }
             StmtKind::Assign { targets, .. } => {
                 if targets.iter().any(|target| self.name_matches(target)) {
                     self.overridden = true;
@@ -150,10 +236,25 @@ where
             visitor::walk_expr(self, expr);
             self.nested = false;
         } else if self.name_matches(expr) {
-            self.usage_count += 1;
+            // If the stack isn't empty, then we're in one of the branches of
+            // a mutually exclusive statement. Otherwise, we'll add it to the
+            // global count.
+            if let Some(last) = self.counter_stack.last_mut() {
+                *last.last_mut().unwrap() += 1;
+            } else {
+                self.usage_count += 1;
+            }
+
+            let current_usage_count = self.usage_count
+                + self
+                    .counter_stack
+                    .iter()
+                    .map(|count| count.last().unwrap_or(&0))
+                    .sum::<u32>();
+
             // For nested loops, the variable usage could be once but the
             // loop makes it being used multiple times.
-            if self.nested || self.usage_count > 1 {
+            if self.nested || current_usage_count > 1 {
                 self.exprs.push(expr);
             }
         } else {

crates/ruff/src/rules/flake8_bugbear/snapshots/ruff__rules__flake8_bugbear__tests__B031_B031.py.snap

@@ -100,4 +100,144 @@ expression: diagnostics
   fix:
     edits: []
   parent: ~
+- kind:
+    name: ReuseOfGroupbyGenerator
+    body: "Using the generator returned from `itertools.groupby()` more than once will do nothing on the second usage"
+    suggestion: ~
+    fixable: false
+  location:
+    row: 86
+    column: 40
+  end_location:
+    row: 86
+    column: 53
+  fix:
+    edits: []
+  parent: ~
+- kind:
+    name: ReuseOfGroupbyGenerator
+    body: "Using the generator returned from `itertools.groupby()` more than once will do nothing on the second usage"
+    suggestion: ~
+    fixable: false
+  location:
+    row: 88
+    column: 40
+  end_location:
+    row: 88
+    column: 53
+  fix:
+    edits: []
+  parent: ~
+- kind:
+    name: ReuseOfGroupbyGenerator
+    body: "Using the generator returned from `itertools.groupby()` more than once will do nothing on the second usage"
+    suggestion: ~
+    fixable: false
+  location:
+    row: 90
+    column: 40
+  end_location:
+    row: 90
+    column: 53
+  fix:
+    edits: []
+  parent: ~
+- kind:
+    name: ReuseOfGroupbyGenerator
+    body: "Using the generator returned from `itertools.groupby()` more than once will do nothing on the second usage"
+    suggestion: ~
+    fixable: false
+  location:
+    row: 91
+    column: 36
+  end_location:
+    row: 91
+    column: 49
+  fix:
+    edits: []
+  parent: ~
+- kind:
+    name: ReuseOfGroupbyGenerator
+    body: "Using the generator returned from `itertools.groupby()` more than once will do nothing on the second usage"
+    suggestion: ~
+    fixable: false
+  location:
+    row: 98
+    column: 48
+  end_location:
+    row: 98
+    column: 61
+  fix:
+    edits: []
+  parent: ~
+- kind:
+    name: ReuseOfGroupbyGenerator
+    body: "Using the generator returned from `itertools.groupby()` more than once will do nothing on the second usage"
+    suggestion: ~
+    fixable: false
+  location:
+    row: 104
+    column: 32
+  end_location:
+    row: 104
+    column: 45
+  fix:
+    edits: []
+  parent: ~
+- kind:
+    name: ReuseOfGroupbyGenerator
+    body: "Using the generator returned from `itertools.groupby()` more than once will do nothing on the second usage"
+    suggestion: ~
+    fixable: false
+  location:
+    row: 113
+    column: 48
+  end_location:
+    row: 113
+    column: 61
+  fix:
+    edits: []
+  parent: ~
+- kind:
+    name: ReuseOfGroupbyGenerator
+    body: "Using the generator returned from `itertools.groupby()` more than once will do nothing on the second usage"
+    suggestion: ~
+    fixable: false
+  location:
+    row: 115
+    column: 48
+  end_location:
+    row: 115
+    column: 61
+  fix:
+    edits: []
+  parent: ~
+- kind:
+    name: ReuseOfGroupbyGenerator
+    body: "Using the generator returned from `itertools.groupby()` more than once will do nothing on the second usage"
+    suggestion: ~
+    fixable: false
+  location:
+    row: 118
+    column: 40
+  end_location:
+    row: 118
+    column: 53
+  fix:
+    edits: []
+  parent: ~
+- kind:
+    name: ReuseOfGroupbyGenerator
+    body: "Using the generator returned from `itertools.groupby()` more than once will do nothing on the second usage"
+    suggestion: ~
+    fixable: false
+  location:
+    row: 122
+    column: 32
+  end_location:
+    row: 122
+    column: 45
+  fix:
+    edits: []
+  parent: ~