From 2cfde403dd187ecee272f8cc59589fb3b43a54ff Mon Sep 17 00:00:00 2001
From: Martin Imre <martinimre25@gmail.com>
Date: Thu, 8 Feb 2024 20:19:16 +0100
Subject: [PATCH] feat(b909): Ignore mutations followed by unconfiditional
 break

---
 bugbear.py            | 26 ++++++++++++++++++--------
 tests/b909.py         | 26 +++++++++++++++++++++++++-
 tests/test_bugbear.py |  2 +-
 3 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/bugbear.py b/bugbear.py
index 1fc7e1d..6663cc4 100644
--- a/bugbear.py
+++ b/bugbear.py
@@ -8,7 +8,7 @@
 import re
 import sys
 import warnings
-from collections import namedtuple
+from collections import defaultdict, namedtuple
 from contextlib import suppress
 from functools import lru_cache, partial
 from keyword import iskeyword
@@ -1583,7 +1583,9 @@ def check_for_b909(self, node: ast.For):
             return
         checker = B909Checker(name)
         checker.visit(node.body)
-        for mutation in checker.mutations:
+        for mutation in itertools.chain.from_iterable(
+            m for m in checker.mutations.values()
+        ):
             self.errors.append(B909(mutation.lineno, mutation.col_offset))
 
 
@@ -1620,17 +1622,18 @@ class B909Checker(ast.NodeVisitor):
 
     def __init__(self, name: str):
         self.name = name
-        self.mutations = []
+        self.mutations = defaultdict(list)
+        self._conditional_block = 0
 
     def visit_Assign(self, node: ast.Assign):
         for target in node.targets:
             if isinstance(target, ast.Subscript) and _to_name_str(target.value):
-                self.mutations.append(node)
+                self.mutations[self._conditional_block].append(node)
         self.generic_visit(node)
 
     def visit_AugAssign(self, node: ast.AugAssign):
         if _to_name_str(node.target) == self.name:
-            self.mutations.append(node)
+            self.mutations[self._conditional_block].append(node)
         self.generic_visit(node)
 
     def visit_Delete(self, node: ast.Delete):
@@ -1644,7 +1647,7 @@ def visit_Delete(self, node: ast.Delete):
                 self.generic_visit(target)
 
             if name == self.name:
-                self.mutations.append(node)
+                self.mutations[self._conditional_block].append(node)
 
     def visit_Call(self, node: ast.Call):
         if isinstance(node.func, ast.Attribute):
@@ -1656,17 +1659,24 @@ def visit_Call(self, node: ast.Call):
                 function_object == self.name
                 and function_name in self.MUTATING_FUNCTIONS
             ):
-                self.mutations.append(node)
+                self.mutations[self._conditional_block].append(node)
 
         self.generic_visit(node)
 
+    def visit_If(self, node: ast.If):
+        self._conditional_block += 1
+        self.visit(node.body)
+        self._conditional_block += 1
+
     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)
+            if isinstance(elem, ast.Break):
+                self.mutations[self._conditional_block].clear()
+            self.visit(elem)
         return node
 
 
diff --git a/tests/b909.py b/tests/b909.py
index 06f35c6..00ee343 100644
--- a/tests/b909.py
+++ b/tests/b909.py
@@ -57,7 +57,7 @@
 
 myset = {1, 2, 3}
 
-for elem in myset:
+for _ in myset:
     # errors
     myset.update({4, 5})
     myset.intersection_update({4, 5})
@@ -103,3 +103,27 @@ def __init__(self, ls):
     foo &= bar
     foo -= bar
     foo ^= bar
+
+
+# more tests for unconditional breaks
+for _ in foo:
+    foo.remove(1)
+    for _ in bar:
+        bar.remove(1)
+        break
+    break
+
+# should not error
+for _ in foo:
+    foo.remove(1)
+    for _ in bar:
+        ...
+    break
+
+# should error (?)
+for _ in foo:
+    foo.remove(1)
+    if bar:
+        bar.remove(1)
+        break
+    break
diff --git a/tests/test_bugbear.py b/tests/test_bugbear.py
index fc3d7cb..53864a6 100644
--- a/tests/test_bugbear.py
+++ b/tests/test_bugbear.py
@@ -974,7 +974,6 @@ def test_b909(self):
         mock_options = Namespace(select=[], extend_select=["B909"])
         bbc = BugBearChecker(filename=str(filename), options=mock_options)
         errors = list(bbc.run())
-        print(errors)
         expected = [
             B909(12, 4),
             B909(13, 4),
@@ -1007,6 +1006,7 @@ def test_b909(self):
             B909(103, 4),
             B909(104, 4),
             B909(105, 4),
+            B909(125, 4),
         ]
         self.assertEqual(errors, self.errors(*expected))