From 3c832e044ae26a2f40820a89f074a3dee7052400 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Tue, 4 Apr 2023 08:37:35 +0530
Subject: [PATCH] fixup! Support mutually exclusive branches for `B031` (#3844)

---
 .../test/fixtures/flake8_bugbear/B031.py      |  7 +++
 .../rules/reuse_of_groupby_generator.rs       |  7 ++-
 ...__flake8_bugbear__tests__B031_B031.py.snap | 52 ++++++++++++-------
 3 files changed, 43 insertions(+), 23 deletions(-)

diff --git a/crates/ruff/resources/test/fixtures/flake8_bugbear/B031.py b/crates/ruff/resources/test/fixtures/flake8_bugbear/B031.py
index 130aba8781926..b034a649acaf9 100644
--- a/crates/ruff/resources/test/fixtures/flake8_bugbear/B031.py
+++ b/crates/ruff/resources/test/fixtures/flake8_bugbear/B031.py
@@ -78,6 +78,13 @@ 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]):
+    if _section == "greens":
+        collect_shop_items(shopper, section_items)
+    else:
+        collect_shop_items(shopper, section_items)
+        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":
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 3aaa47f96a3c5..c2c2bf50a39a4 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
@@ -153,10 +153,9 @@ where
                     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 { .. }));
+                let has_else = orelse
+                    .first()
+                    .map_or(false, |expr| !matches!(expr.node, StmtKind::If { .. }));
 
                 self.parent_ifs.push(stmt);
                 if has_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 f02759d833619..e1248c0954027 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
@@ -107,9 +107,23 @@ expression: diagnostics
     fixable: false
   location:
     row: 86
-    column: 40
+    column: 36
   end_location:
     row: 86
+    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: 93
+    column: 40
+  end_location:
+    row: 93
     column: 53
   fix:
     edits: []
@@ -120,10 +134,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 88
+    row: 95
     column: 40
   end_location:
-    row: 88
+    row: 95
     column: 53
   fix:
     edits: []
@@ -134,10 +148,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 90
+    row: 97
     column: 40
   end_location:
-    row: 90
+    row: 97
     column: 53
   fix:
     edits: []
@@ -148,10 +162,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 91
+    row: 98
     column: 36
   end_location:
-    row: 91
+    row: 98
     column: 49
   fix:
     edits: []
@@ -162,10 +176,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 98
+    row: 105
     column: 48
   end_location:
-    row: 98
+    row: 105
     column: 61
   fix:
     edits: []
@@ -176,10 +190,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 104
+    row: 111
     column: 32
   end_location:
-    row: 104
+    row: 111
     column: 45
   fix:
     edits: []
@@ -190,10 +204,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 113
+    row: 120
     column: 48
   end_location:
-    row: 113
+    row: 120
     column: 61
   fix:
     edits: []
@@ -204,10 +218,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 115
+    row: 122
     column: 48
   end_location:
-    row: 115
+    row: 122
     column: 61
   fix:
     edits: []
@@ -218,10 +232,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 118
+    row: 125
     column: 40
   end_location:
-    row: 118
+    row: 125
     column: 53
   fix:
     edits: []
@@ -232,10 +246,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 122
+    row: 129
     column: 32
   end_location:
-    row: 122
+    row: 129
     column: 45
   fix:
     edits: []