B020: Fix comprehension false postives (#238)

Author: jpy-git

bugbear.py

@@ -557,7 +557,7 @@ def check_for_b020(self, node):
         targets.visit(node.target)
         ctrl_names = set(targets.names)
 
-        iterset = NameFinder()
+        iterset = B020NameFinder()
         iterset.visit(node.iter)
         iterset_names = set(iterset.names)
 
@@ -778,6 +778,22 @@ def visit(self, node):
         return node
 
 
+class B020NameFinder(NameFinder):
+    """Ignore names defined within the local scope of a comprehension."""
+
+    def visit_GeneratorExp(self, node):
+        self.visit(node.generators)
+
+    def visit_ListComp(self, node):
+        self.visit(node.generators)
+
+    def visit_DictComp(self, node):
+        self.visit(node.generators)
+
+    def visit_comprehension(self, node):
+        self.visit(node.iter)
+
+
 error = namedtuple("error", "lineno col message type vars")
 Error = partial(partial, error, type=BugBearChecker, vars=())
 

tests/b020.py

@@ -1,6 +1,6 @@
 """
 Should emit:
-B020 - on lines 8 and 21
+B020 - on lines 8, 21, and 36
 """
 
 items = [1, 2, 3]
@@ -20,3 +20,18 @@
 
 for key, values in values.items():
     print(f"{key}, {values}")
+
+# Variables defined in a comprehension are local in scope
+# to that comprehension and are therefore allowed.
+for var in [var for var in range(10)]:
+    print(var)
+
+for var in (var for var in range(10)):
+    print(var)
+
+for k, v in {k: v for k, v in zip(range(10), range(10, 20))}.items():
+    print(k, v)
+
+# However we still call out reassigning the iterable in the comprehension.
+for vars in [i for i in vars]:
+    print(vars)

tests/test_bugbear.py

@@ -284,6 +284,7 @@ def test_b020(self):
             self.errors(
                 B020(8, 4, vars=("items",)),
                 B020(21, 9, vars=("values",)),
+                B020(36, 4, vars=("vars",)),
             ),
         )