diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B031.py b/crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B031.py index 61f62c4b67..412b890ef7 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B031.py +++ b/crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B031.py @@ -174,6 +174,49 @@ for (_key1, _key2), (_value1, _value2) in groupby( collect_shop_items("Jane", group[1]) collect_shop_items("Joe", group[1]) +# Shouldn't trigger the warning when there is a continue, break statement. +for _section, section_items in groupby(items, key=lambda p: p[1]): + if _section == "greens": + collect_shop_items(shopper, section_items) + continue + elif _section == "frozen items": + collect_shop_items(shopper, section_items) + break + collect_shop_items(shopper, section_items) + +# Shouldn't trigger the warning when there is a return statement. +for _section, section_items in groupby(items, key=lambda p: p[1]): + if _section == "greens": + collect_shop_items(shopper, section_items) + return + elif _section == "frozen items": + return section_items + collect_shop_items(shopper, section_items) + +# Should trigger the warning for duplicate access, even if is a return statement after. +for _section, section_items in groupby(items, key=lambda p: p[1]): + if _section == "greens": + collect_shop_items(shopper, section_items) + collect_shop_items(shopper, section_items) + return + +# Should trigger the warning for duplicate access, even if is a return in another branch. +for _section, section_items in groupby(items, key=lambda p: p[1]): + if _section == "greens": + collect_shop_items(shopper, section_items) + return + elif _section == "frozen items": + collect_shop_items(shopper, section_items) + collect_shop_items(shopper, section_items) + +# Should trigger, since only one branch has a return statement. +for _section, section_items in groupby(items, key=lambda p: p[1]): + if _section == "greens": + collect_shop_items(shopper, section_items) + return + elif _section == "frozen items": + collect_shop_items(shopper, section_items) + collect_shop_items(shopper, section_items) # B031 # Let's redefine the `groupby` function to make sure we pick up the correct one. # NOTE: This should always be at the end of the file. diff --git a/crates/ruff_linter/src/rules/flake8_bugbear/rules/reuse_of_groupby_generator.rs b/crates/ruff_linter/src/rules/flake8_bugbear/rules/reuse_of_groupby_generator.rs index 2b39e349af..b495276c02 100644 --- a/crates/ruff_linter/src/rules/flake8_bugbear/rules/reuse_of_groupby_generator.rs +++ b/crates/ruff_linter/src/rules/flake8_bugbear/rules/reuse_of_groupby_generator.rs @@ -100,6 +100,16 @@ impl<'a> GroupNameFinder<'a> { self.usage_count += value; } } + + /// Reset the usage count for the group name by the given value. + /// This function is called when there is a `continue`, `break`, or `return` statement. + fn reset_usage_count(&mut self) { + if let Some(last) = self.counter_stack.last_mut() { + *last.last_mut().unwrap() = 0; + } else { + self.usage_count = 0; + } + } } impl<'a> Visitor<'a> for GroupNameFinder<'a> { @@ -197,6 +207,15 @@ impl<'a> Visitor<'a> for GroupNameFinder<'a> { self.visit_expr(expr); } } + Stmt::Continue(_) | Stmt::Break(_) => { + self.reset_usage_count(); + } + Stmt::Return(ast::StmtReturn { value, range: _ }) => { + if let Some(expr) = value { + self.visit_expr(expr); + } + self.reset_usage_count(); + } _ => visitor::walk_stmt(self, stmt), } } diff --git a/crates/ruff_linter/src/rules/flake8_bugbear/snapshots/ruff_linter__rules__flake8_bugbear__tests__B031_B031.py.snap b/crates/ruff_linter/src/rules/flake8_bugbear/snapshots/ruff_linter__rules__flake8_bugbear__tests__B031_B031.py.snap index b2de85a2f4..075bffbe6c 100644 --- a/crates/ruff_linter/src/rules/flake8_bugbear/snapshots/ruff_linter__rules__flake8_bugbear__tests__B031_B031.py.snap +++ b/crates/ruff_linter/src/rules/flake8_bugbear/snapshots/ruff_linter__rules__flake8_bugbear__tests__B031_B031.py.snap @@ -195,4 +195,31 @@ B031.py:144:33: B031 Using the generator returned from `itertools.groupby()` mor 146 | for group in groupby(items, key=lambda p: p[1]): | +B031.py:200:37: B031 Using the generator returned from `itertools.groupby()` more than once will do nothing on the second usage + | +198 | if _section == "greens": +199 | collect_shop_items(shopper, section_items) +200 | collect_shop_items(shopper, section_items) + | ^^^^^^^^^^^^^ B031 +201 | return + | +B031.py:210:37: B031 Using the generator returned from `itertools.groupby()` more than once will do nothing on the second usage + | +208 | elif _section == "frozen items": +209 | collect_shop_items(shopper, section_items) +210 | collect_shop_items(shopper, section_items) + | ^^^^^^^^^^^^^ B031 +211 | +212 | # Should trigger, since only one branch has a return statement. + | + +B031.py:219:33: B031 Using the generator returned from `itertools.groupby()` more than once will do nothing on the second usage + | +217 | elif _section == "frozen items": +218 | collect_shop_items(shopper, section_items) +219 | collect_shop_items(shopper, section_items) # B031 + | ^^^^^^^^^^^^^ B031 +220 | +221 | # Let's redefine the `groupby` function to make sure we pick up the correct one. + |