Visit comprehension to detect group name usage/overrides (#3887)

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

@@ -78,6 +78,21 @@ for _section, section_items in itertools.groupby(items, key=lambda p: p[1]):
     for shopper in shoppers:
         collect_shop_items(shopper, section_items)  # B031
 
+for _section, section_items in itertools.groupby(items, key=lambda p: p[1]):
+    _ = [collect_shop_items(shopper, section_items) for shopper in shoppers]  # B031
+
+for _section, section_items in itertools.groupby(items, key=lambda p: p[1]):
+    # The variable is overridden, skip checking.
+    _ = [_ for section_items in range(3)]
+    _ = [collect_shop_items(shopper, section_items) for shopper in shoppers]
+
+for _section, section_items in itertools.groupby(items, key=lambda p: p[1]):
+    _ = [item for item in section_items]
+
+for _section, section_items in itertools.groupby(items, key=lambda p: p[1]):
+    # The iterator is being used for the second time.
+    _ = [(item1, item2) for item1 in section_items for item2 in 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)

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

@@ -1,4 +1,4 @@
-use rustpython_parser::ast::{Expr, ExprKind, Stmt, StmtKind};
+use rustpython_parser::ast::{Comprehension, Expr, ExprKind, Stmt, StmtKind};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
@@ -204,20 +204,39 @@ where
                     }
                 }
             }
-            StmtKind::Assign { targets, .. } => {
+            StmtKind::Assign { targets, value, .. } => {
                 if targets.iter().any(|target| self.name_matches(target)) {
                     self.overridden = true;
+                } else {
+                    self.visit_expr(value);
                 }
             }
-            StmtKind::AnnAssign { target, .. } => {
+            StmtKind::AnnAssign { target, value, .. } => {
                 if self.name_matches(target) {
                     self.overridden = true;
+                } else if let Some(expr) = value {
+                    self.visit_expr(expr);
                 }
             }
             _ => visitor::walk_stmt(self, stmt),
         }
     }
 
+    fn visit_comprehension(&mut self, comprehension: &'a Comprehension) {
+        if self.name_matches(&comprehension.target) {
+            self.overridden = true;
+        }
+        if self.overridden {
+            return;
+        }
+        if self.name_matches(&comprehension.iter) {
+            self.usage_count += 1;
+            if self.usage_count > 1 {
+                self.exprs.push(&comprehension.iter);
+            }
+        }
+    }
+
     fn visit_expr(&mut self, expr: &'a Expr) {
         if let ExprKind::NamedExpr { target, .. } = &expr.node {
             if self.name_matches(target) {
@@ -227,14 +246,35 @@ where
         if self.overridden {
             return;
         }
-        if matches!(
-            &expr.node,
-            ExprKind::ListComp { .. } | ExprKind::DictComp { .. } | ExprKind::SetComp { .. }
-        ) {
+
+        match &expr.node {
+            ExprKind::ListComp { elt, generators } | ExprKind::SetComp { elt, generators } => {
+                for comprehension in generators {
+                    self.visit_comprehension(comprehension);
+                }
+                if !self.overridden {
                     self.nested = true;
-            visitor::walk_expr(self, expr);
+                    visitor::walk_expr(self, elt);
                     self.nested = false;
-        } else if self.name_matches(expr) {
+                }
+            }
+            ExprKind::DictComp {
+                key,
+                value,
+                generators,
+            } => {
+                for comprehension in generators {
+                    self.visit_comprehension(comprehension);
+                }
+                if !self.overridden {
+                    self.nested = true;
+                    visitor::walk_expr(self, key);
+                    visitor::walk_expr(self, value);
+                    self.nested = false;
+                }
+            }
+            _ => {
+                if self.name_matches(expr) {
                     // 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.
@@ -261,6 +301,8 @@ where
                 }
             }
         }
+    }
+}
 
 /// B031
 pub fn reuse_of_groupby_generator(

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

@@ -106,10 +106,38 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 86
+    row: 82
+    column: 37
+  end_location:
+    row: 82
+    column: 50
+  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: 94
+    column: 64
+  end_location:
+    row: 94
+    column: 77
+  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: 101
     column: 36
   end_location:
-    row: 86
+    row: 101
     column: 49
   fix:
     edits: []
@@ -120,10 +148,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 93
+    row: 108
     column: 40
   end_location:
-    row: 93
+    row: 108
     column: 53
   fix:
     edits: []
@@ -134,10 +162,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 95
+    row: 110
     column: 40
   end_location:
-    row: 95
+    row: 110
     column: 53
   fix:
     edits: []
@@ -148,10 +176,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 97
+    row: 112
     column: 40
   end_location:
-    row: 97
+    row: 112
     column: 53
   fix:
     edits: []
@@ -162,10 +190,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 98
+    row: 113
     column: 36
   end_location:
-    row: 98
+    row: 113
     column: 49
   fix:
     edits: []
@@ -176,10 +204,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 105
+    row: 120
     column: 48
   end_location:
-    row: 105
+    row: 120
     column: 61
   fix:
     edits: []
@@ -190,10 +218,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 111
+    row: 126
     column: 32
   end_location:
-    row: 111
+    row: 126
     column: 45
   fix:
     edits: []
@@ -204,10 +232,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 120
+    row: 135
     column: 48
   end_location:
-    row: 120
+    row: 135
     column: 61
   fix:
     edits: []
@@ -218,10 +246,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 122
+    row: 137
     column: 48
   end_location:
-    row: 122
+    row: 137
     column: 61
   fix:
     edits: []
@@ -232,10 +260,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 125
+    row: 140
     column: 40
   end_location:
-    row: 125
+    row: 140
     column: 53
   fix:
     edits: []
@@ -246,10 +274,10 @@ expression: diagnostics
     suggestion: ~
     fixable: false
   location:
-    row: 129
+    row: 144
     column: 32
   end_location:
-    row: 129
+    row: 144
     column: 45
   fix:
     edits: []