From 88da797df3f1cc9d26a1df896ca017ad8624ef7c Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Wed, 17 Jul 2024 10:21:17 -0400
Subject: [PATCH] Detect enumerate iterations in loop-iterator-mutation

---
 .../test/fixtures/flake8_bugbear/B909.py      |  4 +++
 .../rules/loop_iterator_mutation.rs           | 26 ++++++++++++++++---
 ...__flake8_bugbear__tests__B909_B909.py.snap |  8 ++++++
 3 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B909.py b/crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B909.py
index 68afaf87fb257b..99349f4a81f494 100644
--- a/crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B909.py
+++ b/crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B909.py
@@ -158,3 +158,7 @@ def __init__(self, ls):
     some_list[elem] = 1
     some_list.remove(elem)
     some_list.discard(elem)
+
+# should error
+for i, elem in enumerate(some_list):
+    some_list.pop(0)
diff --git a/crates/ruff_linter/src/rules/flake8_bugbear/rules/loop_iterator_mutation.rs b/crates/ruff_linter/src/rules/flake8_bugbear/rules/loop_iterator_mutation.rs
index d48466813acb7d..847f78dcb6bc8e 100644
--- a/crates/ruff_linter/src/rules/flake8_bugbear/rules/loop_iterator_mutation.rs
+++ b/crates/ruff_linter/src/rules/flake8_bugbear/rules/loop_iterator_mutation.rs
@@ -64,9 +64,29 @@ pub(crate) fn loop_iterator_mutation(checker: &mut Checker, stmt_for: &StmtFor)
         range: _,
     } = stmt_for;
 
-    if !matches!(iter.as_ref(), Expr::Name(_) | Expr::Attribute(_)) {
-        return;
-    }
+    let (target, iter) = match iter.as_ref() {
+        Expr::Name(_) | Expr::Attribute(_) => {
+            // Ex) `for item in items:`
+            (&**target, &**iter)
+        }
+        Expr::Call(ExprCall {
+            func, arguments, ..
+        }) => {
+            // Ex) `for i, item in enumerate(items):`
+            if checker.semantic().match_builtin_expr(func, "enumerate") {
+                if let [iter] = &*arguments.args {
+                    (&**target, iter)
+                } else {
+                    return;
+                }
+            } else {
+                return;
+            }
+        }
+        _ => {
+            return;
+        }
+    };
 
     // Collect mutations to the iterable.
     let mutations = {
diff --git a/crates/ruff_linter/src/rules/flake8_bugbear/snapshots/ruff_linter__rules__flake8_bugbear__tests__B909_B909.py.snap b/crates/ruff_linter/src/rules/flake8_bugbear/snapshots/ruff_linter__rules__flake8_bugbear__tests__B909_B909.py.snap
index 7f70841c6c0664..80e740050197e2 100644
--- a/crates/ruff_linter/src/rules/flake8_bugbear/snapshots/ruff_linter__rules__flake8_bugbear__tests__B909_B909.py.snap
+++ b/crates/ruff_linter/src/rules/flake8_bugbear/snapshots/ruff_linter__rules__flake8_bugbear__tests__B909_B909.py.snap
@@ -339,3 +339,11 @@ B909.py:150:8: B909 Mutation to loop iterable `some_list` during iteration
 151 |         pass
 152 |     else:
     |
+
+B909.py:164:5: B909 Mutation to loop iterable `some_list` during iteration
+    |
+162 | # should error
+163 | for i, elem in enumerate(some_list):
+164 |     some_list.pop(0)
+    |     ^^^^^^^^^^^^^ B909
+    |