diff --git a/crates/ruff/resources/test/fixtures/flake8_bugbear/B031.py b/crates/ruff/resources/test/fixtures/flake8_bugbear/B031.py index b034a649ac..1e8693ae17 100644 --- a/crates/ruff/resources/test/fixtures/flake8_bugbear/B031.py +++ b/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) 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 c2c2bf50a3..c456e8af9a 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 @@ -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,37 +246,60 @@ where if self.overridden { return; } - if matches!( - &expr.node, - ExprKind::ListComp { .. } | ExprKind::DictComp { .. } | ExprKind::SetComp { .. } - ) { - self.nested = true; - visitor::walk_expr(self, expr); - self.nested = false; - } else 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. - 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 || current_usage_count > 1 { - self.exprs.push(expr); + 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, elt); + self.nested = false; + } + } + 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. + 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 || current_usage_count > 1 { + self.exprs.push(expr); + } + } else { + visitor::walk_expr(self, expr); + } } - } else { - visitor::walk_expr(self, expr); } } } 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 e1248c0954..dbf4fabc89 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 @@ -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: []