[flake8-bugbear] Avoid false positive for usage after continue (B031)

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

@@ -174,6 +174,49 @@ def collect_shop_items(shopper, items):
     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.

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),
         }
     }

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.
+    |