From fd4e612377ed288fe530bae39404930c37617248 Mon Sep 17 00:00:00 2001
From: Martin Imre <martinimre25@gmail.com>
Date: Fri, 12 Jan 2024 19:42:48 +0100
Subject: [PATCH] feat(rules): Add rule to check for mutations of loop iterator

---
 bugbear.py            | 59 +++++++++++++++++++++++++++++++++++++++++++
 tests/b999.py         | 46 +++++++++++++++++++++++++++++++++
 tests/test_bugbear.py | 16 ++++++++++++
 3 files changed, 121 insertions(+)
 create mode 100644 tests/b999.py

diff --git a/bugbear.py b/bugbear.py
index 1e90dc6..61e5a0c 100644
--- a/bugbear.py
+++ b/bugbear.py
@@ -235,6 +235,7 @@ def _flatten_excepthandler(node):
         ):
             expr_list.extend(expr.value.elts)
             continue
+
         yield expr
 
 
@@ -495,6 +496,7 @@ def visit_For(self, node):
         self.check_for_b020(node)
         self.check_for_b023(node)
         self.check_for_b031(node)
+        self.check_for_b999(node)
         self.generic_visit(node)
 
     def visit_AsyncFor(self, node):
@@ -1544,6 +1546,18 @@ def check(num_args, param_name):
         elif node.func.attr == "split":
             check(2, "maxsplit")
 
+    def check_for_b999(self, node: ast.For):
+        if isinstance(node.iter, ast.Name):
+            name = _to_name_str(node.iter)
+        elif isinstance(node.iter, ast.Attribute):
+            name = _to_name_str(node.iter)
+        else:
+            return
+        checker = B999Checker(name)
+        checker.visit(node.body)
+        for mutation in checker.mutations:
+            self.errors.append(B999(mutation.lineno, mutation.col_offset))
+
 
 def compose_call_path(node):
     if isinstance(node, ast.Attribute):
@@ -1555,6 +1569,49 @@ def compose_call_path(node):
         yield node.id
 
 
+class B999Checker(ast.NodeVisitor):
+    def __init__(self, name: str):
+        self.name = name
+        self.mutations = []
+
+    def visit_Delete(self, node: ast.Delete):
+        for target in node.targets:
+            if isinstance(target, ast.Subscript):
+                name = _to_name_str(target.value)
+            elif isinstance(target, (ast.Attribute, ast.Name)):
+                name = _to_name_str(target)
+            else:
+                name = ""  # fallback
+                self.generic_visit(target)
+
+            if name == self.name:
+                self.mutations.append(node)
+
+    def visit_Call(self, node: ast.Call):
+        if isinstance(node.func, ast.Attribute):
+            name = _to_name_str(node.func.value)
+            function_object = name
+
+            if function_object == self.name:
+                self.mutations.append(node)
+
+        self.generic_visit(node)
+
+    def visit_Name(self, node: ast.Name):
+        if isinstance(node.ctx, ast.Del):
+            self.mutations.append(node)
+        self.generic_visit(node)
+
+    def visit(self, node):
+        """Like super-visit but supports iteration over lists."""
+        if not isinstance(node, list):
+            return super().visit(node)
+
+        for elem in node:
+            super().visit(elem)
+        return node
+
+
 @attr.s
 class NameFinder(ast.NodeVisitor):
     """Finds a name within a tree of nodes.
@@ -2045,6 +2102,8 @@ def visit_Lambda(self, node):
         " statement."
     )
 )
+
 B950 = Error(message="B950 line too long ({} > {} characters)")
 
+B999 = Error(message="B999 editing loop iterable can cause unintended bugs!")
 disabled_by_default = ["B901", "B902", "B903", "B904", "B905", "B906", "B908", "B950"]
diff --git a/tests/b999.py b/tests/b999.py
new file mode 100644
index 0000000..9bec7ba
--- /dev/null
+++ b/tests/b999.py
@@ -0,0 +1,46 @@
+"""
+Should emit:
+B999 - on lines 11, 25, 26, 40, 46
+"""
+
+
+some_list = [1, 2, 3]
+for elem in some_list:
+    print(elem)
+    if elem % 2 == 0:
+        some_list.remove(elem)  # should error
+
+some_list = [1, 2, 3]
+some_other_list = [1, 2, 3]
+for elem in some_list:
+    print(elem)
+    if elem % 2 == 0:
+        some_other_list.remove(elem)  # should not error
+
+
+some_list = [1, 2, 3]
+for elem in some_list:
+    print(elem)
+    if elem % 2 == 0:
+        del some_list[2]  # should error
+        del some_list
+
+
+class A:
+    some_list: list
+
+    def __init__(self, ls):
+        self.some_list = list(ls)
+
+
+a = A((1, 2, 3))
+for elem in a.some_list:
+    print(elem)
+    if elem % 2 == 0:
+        a.some_list.remove(elem)  # should error
+
+a = A((1, 2, 3))
+for elem in a.some_list:
+    print(elem)
+    if elem % 2 == 0:
+        del a.some_list[2]  # should error
diff --git a/tests/test_bugbear.py b/tests/test_bugbear.py
index 08f106b..84e5ce2 100644
--- a/tests/test_bugbear.py
+++ b/tests/test_bugbear.py
@@ -54,6 +54,7 @@
     B907,
     B908,
     B950,
+    B999,
     BugBearChecker,
     BugBearVisitor,
 )
@@ -952,6 +953,21 @@ def test_selfclean_test_bugbear(self):
         self.assertEqual(proc.stdout, b"")
         self.assertEqual(proc.stderr, b"")
 
+    def test_b999(self):
+        filename = Path(__file__).absolute().parent / "b999.py"
+        mock_options = Namespace(select=[], extend_select=["B999"])
+        bbc = BugBearChecker(filename=str(filename), options=mock_options)
+        errors = list(bbc.run())
+        print(errors)
+        expected = [
+            B999(11, 8),
+            B999(25, 8),
+            B999(26, 8),
+            B999(40, 8),
+            B999(46, 8),
+        ]
+        self.assertEqual(errors, self.errors(*expected))
+
 
 class TestFuzz(unittest.TestCase):
     # TODO: enable this test on py312 once hypothesmith supports py312