diff --git a/crates/ruff/resources/test/fixtures/flake8_bugbear/B031.py b/crates/ruff/resources/test/fixtures/flake8_bugbear/B031.py index 82a2c0a3f9888..130aba8781926 100644 --- a/crates/ruff/resources/test/fixtures/flake8_bugbear/B031.py +++ b/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]) diff --git a/crates/ruff/src/rules/flake8_bugbear/rules/reuse_of_groupby_generator.rs b/crates/ruff/src/rules/flake8_bugbear/rules/reuse_of_groupby_generator.rs index 5623980955e30..3aaa47f96a3c5 100644 --- a/crates/ruff/src/rules/flake8_bugbear/rules/reuse_of_groupby_generator.rs +++ b/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>, /// 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::(); + // 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 { diff --git a/crates/ruff/src/rules/flake8_bugbear/snapshots/ruff__rules__flake8_bugbear__tests__B031_B031.py.snap b/crates/ruff/src/rules/flake8_bugbear/snapshots/ruff__rules__flake8_bugbear__tests__B031_B031.py.snap index 320d1c7f8acb6..f02759d833619 100644 --- a/crates/ruff/src/rules/flake8_bugbear/snapshots/ruff__rules__flake8_bugbear__tests__B031_B031.py.snap +++ b/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: ~