diff --git a/bugbear.py b/bugbear.py index 542bece..d1f5cd8 100644 --- a/bugbear.py +++ b/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=()) diff --git a/tests/b020.py b/tests/b020.py index df30e75..4f30984 100644 --- a/tests/b020.py +++ b/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) diff --git a/tests/test_bugbear.py b/tests/test_bugbear.py index 6639f17..7954b5d 100644 --- a/tests/test_bugbear.py +++ b/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",)), ), )