feat(rules): Add rule to check for mutations of loop iterator (#446)

* feat(rules): Add rule to check for mutations of loop iterator

* fixup! feat(rules): Add rule to check for mutations of loop iterator

* doc(B038): Add doc for B038 to README.rst

---------

Co-authored-by: Cooper Lees <[email protected]>

Author: mimre25

README.rst

@@ -198,6 +198,10 @@ second usage. Save the result to a list if the result is needed multiple times.
 **B036**: Found ``except BaseException:`` without re-raising (no ``raise`` in the top-level of the ``except`` block). This catches all kinds of things (Exception, SystemExit, KeyboardInterrupt...) and may prevent a program from exiting as expected.
 
 **B037**: Found ``return <value>``, ``yield``, ``yield <value>``, or ``yield from <value>`` in class ``__init__()`` method. No values should be returned or yielded, only bare ``return``s are ok.
+
+**B038**: Found a mutation of a mutable loop iterable inside the loop body.
+Changes to the iterable of a loop such as calls to `list.remove()` or via `del` can cause unintended bugs.
+
 Opinionated warnings
 ~~~~~~~~~~~~~~~~~~~~
 

bugbear.py

@@ -237,6 +237,7 @@ def _flatten_excepthandler(node):
         ):
             expr_list.extend(expr.value.elts)
             continue
+
         yield expr
 
 
@@ -521,6 +522,7 @@ def visit_For(self, node):
         self.check_for_b020(node)
         self.check_for_b023(node)
         self.check_for_b031(node)
+        self.check_for_b038(node)
         self.generic_visit(node)
 
     def visit_AsyncFor(self, node):
@@ -1570,6 +1572,18 @@ def check(num_args, param_name):
         elif node.func.attr == "split":
             check(2, "maxsplit")
 
+    def check_for_b038(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 = B038Checker(name)
+        checker.visit(node.body)
+        for mutation in checker.mutations:
+            self.errors.append(B038(mutation.lineno, mutation.col_offset))
+
 
 def compose_call_path(node):
     if isinstance(node, ast.Attribute):
@@ -1581,6 +1595,49 @@ def compose_call_path(node):
         yield node.id
 
 
+class B038Checker(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.
@@ -2075,6 +2132,12 @@ def visit_Lambda(self, node):
         " statement."
     )
 )
+
 B950 = Error(message="B950 line too long ({} > {} characters)")
 
+B038 = Error(
+    message=(
+        "B038 editing a loop's mutable iterable often leads to unexpected results/bugs"
+    )
+)
 disabled_by_default = ["B901", "B902", "B903", "B904", "B905", "B906", "B908", "B950"]

tests/b038.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

tests/test_bugbear.py

@@ -46,6 +46,7 @@
     B035,
     B036,
     B037,
+    B038,
     B901,
     B902,
     B903,
@@ -967,6 +968,21 @@ def test_selfclean_test_bugbear(self):
         self.assertEqual(proc.stdout, b"")
         self.assertEqual(proc.stderr, b"")
 
+    def test_b038(self):
+        filename = Path(__file__).absolute().parent / "b038.py"
+        mock_options = Namespace(select=[], extend_select=["B038"])
+        bbc = BugBearChecker(filename=str(filename), options=mock_options)
+        errors = list(bbc.run())
+        print(errors)
+        expected = [
+            B038(11, 8),
+            B038(25, 8),
+            B038(26, 8),
+            B038(40, 8),
+            B038(46, 8),
+        ]
+        self.assertEqual(errors, self.errors(*expected))
+
 
 class TestFuzz(unittest.TestCase):
     from hypothesis import HealthCheck, given, settings