Implement late-binding loop check (#265)

* Implement late-binding loop check

* Dedupe warning with nested loops

* Only trigger for Load names

* Handle attribute/item assignment

* Allow explictly-captured vars

i.e. don't complain about the default value of arguments, since that's an explicit solution to late binding!

Author: Zac-HD

README.rst

@@ -150,6 +150,9 @@ No exceptions will be suppressed and therefore this context manager is redundant
 N.B. this rule currently does not flag `suppress` calls to avoid potential false
 positives due to similarly named user-defined functions.
 
+**B023**: Functions defined inside a loop must not use variables redefined in
+the loop, because `late-binding closures are a classic gotcha
+<https://docs.python-guide.org/writing/gotchas/#late-binding-closures>`__.
 
 Opinionated warnings
 ~~~~~~~~~~~~~~~~~~~~

bugbear.py

@@ -26,6 +26,7 @@
     ast.DictComp,
     ast.GeneratorExp,
 )
+FUNCTION_NODES = (ast.AsyncFunctionDef, ast.FunctionDef, ast.Lambda)
 
 Context = namedtuple("Context", ["node", "stack"])
 
@@ -198,6 +199,23 @@ def _to_name_str(node):
         return _to_name_str(node.value)
 
 
+def names_from_assignments(assign_target):
+    if isinstance(assign_target, ast.Name):
+        yield assign_target.id
+    elif isinstance(assign_target, ast.Starred):
+        yield from names_from_assignments(assign_target.value)
+    elif isinstance(assign_target, (ast.List, ast.Tuple)):
+        for child in assign_target.elts:
+            yield from names_from_assignments(child)
+
+
+def children_in_scope(node):
+    yield node
+    if not isinstance(node, FUNCTION_NODES):
+        for child in ast.iter_child_nodes(node):
+            yield from children_in_scope(child)
+
+
 def _typesafe_issubclass(cls, class_or_tuple):
     try:
         return issubclass(cls, class_or_tuple)
@@ -220,6 +238,7 @@ class BugBearVisitor(ast.NodeVisitor):
     contexts = attr.ib(default=attr.Factory(list))
 
     NODE_WINDOW_SIZE = 4
+    _b023_seen = attr.ib(factory=set, init=False)
 
     if False:
         # Useful for tracing what the hell is going on.
@@ -348,6 +367,31 @@ def visit_Assign(self, node):
     def visit_For(self, node):
         self.check_for_b007(node)
         self.check_for_b020(node)
+        self.check_for_b023(node)
+        self.generic_visit(node)
+
+    def visit_AsyncFor(self, node):
+        self.check_for_b023(node)
+        self.generic_visit(node)
+
+    def visit_While(self, node):
+        self.check_for_b023(node)
+        self.generic_visit(node)
+
+    def visit_ListComp(self, node):
+        self.check_for_b023(node)
+        self.generic_visit(node)
+
+    def visit_SetComp(self, node):
+        self.check_for_b023(node)
+        self.generic_visit(node)
+
+    def visit_DictComp(self, node):
+        self.check_for_b023(node)
+        self.generic_visit(node)
+
+    def visit_GeneratorExp(self, node):
+        self.check_for_b023(node)
         self.generic_visit(node)
 
     def visit_Assert(self, node):
@@ -520,6 +564,59 @@ def check_for_b020(self, node):
                 n = targets.names[name][0]
                 self.errors.append(B020(n.lineno, n.col_offset, vars=(name,)))
 
+    def check_for_b023(self, loop_node):
+        """Check that functions (including lambdas) do not use loop variables.
+
+        https://docs.python-guide.org/writing/gotchas/#late-binding-closures from
+        functions - usually but not always lambdas - defined inside a loop are a
+        classic source of bugs.
+
+        For each use of a variable inside a function defined inside a loop, we
+        emit a warning if that variable is reassigned on each loop iteration
+        (outside the function).  This includes but is not limited to explicit
+        loop variables like the `x` in `for x in range(3):`.
+        """
+        # Because most loops don't contain functions, it's most efficient to
+        # implement this "backwards": first we find all the candidate variable
+        # uses, and then if there are any we check for assignment of those names
+        # inside the loop body.
+        suspicious_variables = []
+        for node in ast.walk(loop_node):
+            if isinstance(node, FUNCTION_NODES):
+                argnames = {
+                    arg.arg for arg in ast.walk(node.args) if isinstance(arg, ast.arg)
+                }
+                if isinstance(node, ast.Lambda):
+                    body_nodes = ast.walk(node.body)
+                else:
+                    body_nodes = itertools.chain.from_iterable(map(ast.walk, node.body))
+                for name in body_nodes:
+                    if (
+                        isinstance(name, ast.Name)
+                        and name.id not in argnames
+                        and isinstance(name.ctx, ast.Load)
+                    ):
+                        err = B023(name.lineno, name.col_offset, vars=(name.id,))
+                        if err not in self._b023_seen:
+                            self._b023_seen.add(err)  # dedupe across nested loops
+                            suspicious_variables.append(err)
+
+        if suspicious_variables:
+            reassigned_in_loop = set(self._get_assigned_names(loop_node))
+
+        for err in sorted(suspicious_variables):
+            if reassigned_in_loop.issuperset(err.vars):
+                self.errors.append(err)
+
+    def _get_assigned_names(self, loop_node):
+        loop_targets = (ast.For, ast.AsyncFor, ast.comprehension)
+        for node in children_in_scope(loop_node):
+            if isinstance(node, (ast.Assign)):
+                for child in node.targets:
+                    yield from names_from_assignments(child)
+            if isinstance(node, loop_targets + (ast.AnnAssign, ast.AugAssign)):
+                yield from names_from_assignments(node.target)
+
     def check_for_b904(self, node):
         """Checks `raise` without `from` inside an `except` clause.
 
@@ -1041,6 +1138,8 @@ def visit_Lambda(self, node):
     )
 )
 
+B023 = Error(message="B023 Function definition does not bind loop variable {!r}.")
+
 # Warnings disabled by default.
 B901 = Error(
     message=(

tests/b023.py

@@ -0,0 +1,78 @@
+"""
+Should emit:
+B023 - on lines 12, 13, 16, 28, 29, 30, 31, 40, 42, 50, 51, 52, 53, 61, 68.
+"""
+
+functions = []
+z = 0
+
+for x in range(3):
+    y = x + 1
+    # Subject to late-binding problems
+    functions.append(lambda: x)
+    functions.append(lambda: y)  # not just the loop var
+
+    def f_bad_1():
+        return x
+
+    # Actually OK
+    functions.append(lambda x: x * 2)
+    functions.append(lambda x=x: x)
+    functions.append(lambda: z)  # OK because not assigned in the loop
+
+    def f_ok_1(x):
+        return x * 2
+
+
+def check_inside_functions_too():
+    ls = [lambda: x for x in range(2)]
+    st = {lambda: x for x in range(2)}
+    gn = (lambda: x for x in range(2))
+    dt = {x: lambda: x for x in range(2)}
+
+
+async def pointless_async_iterable():
+    yield 1
+
+
+async def container_for_problems():
+    async for x in pointless_async_iterable():
+        functions.append(lambda: x)
+
+    [lambda: x async for x in pointless_async_iterable()]
+
+
+a = 10
+b = 0
+while True:
+    a = a_ = a - 1
+    b += 1
+    functions.append(lambda: a)
+    functions.append(lambda: a_)
+    functions.append(lambda: b)
+    functions.append(lambda: c)  # not a name error because of late binding!
+    c: bool = a > 3
+    if not c:
+        break
+
+# Nested loops should not duplicate reports
+for j in range(2):
+    for k in range(3):
+        lambda: j * k
+
+
+for j, k, l in [(1, 2, 3)]:
+
+    def f():
+        j = None  # OK because it's an assignment
+        [l for k in range(2)]  # error for l, not for k
+
+        assert a and functions
+
+    a.attribute = 1  # modifying an attribute doesn't make it a loop variable
+    functions[0] = lambda: None  # same for an element
+
+for var in range(2):
+
+    def explicit_capture(captured=var):
+        return captured

tests/test_bugbear.py

@@ -33,6 +33,7 @@
     B020,
     B021,
     B022,
+    B023,
     B901,
     B902,
     B903,
@@ -325,6 +326,30 @@ def test_b022(self):
         errors = list(bbc.run())
         self.assertEqual(errors, self.errors(B022(8, 0)))
 
+    def test_b023(self):
+        filename = Path(__file__).absolute().parent / "b023.py"
+        bbc = BugBearChecker(filename=str(filename))
+        errors = list(bbc.run())
+        expected = self.errors(
+            B023(12, 29, vars=("x",)),
+            B023(13, 29, vars=("y",)),
+            B023(16, 15, vars=("x",)),
+            B023(28, 18, vars=("x",)),
+            B023(29, 18, vars=("x",)),
+            B023(30, 18, vars=("x",)),
+            B023(31, 21, vars=("x",)),
+            B023(40, 33, vars=("x",)),
+            B023(42, 13, vars=("x",)),
+            B023(50, 29, vars=("a",)),
+            B023(51, 29, vars=("a_",)),
+            B023(52, 29, vars=("b",)),
+            B023(53, 29, vars=("c",)),
+            B023(61, 16, vars=("j",)),
+            B023(61, 20, vars=("k",)),
+            B023(68, 9, vars=("l",)),
+        )
+        self.assertEqual(errors, expected)
+
     def test_b901(self):
         filename = Path(__file__).absolute().parent / "b901.py"
         bbc = BugBearChecker(filename=str(filename))