B023: don't flag a function that is only called inside the loop (#567)

* B023: don't flag a function that is only called inside the loop

check_for_b023 warns whenever a function defined in a loop closes over the
loop variable, but a function whose every reference is a direct call in the
loop body cannot outlive the iteration it was defined in, so the value it
closes over is the one the author meant.

The existing safe_functions notion covered only a fixed set of shapes -
filter/map/reduce, a key= argument, and 'return lambda: x'. It is replaced by
the rule the maintainers described on the two issues: warn when the name
escapes the loop or is referenced as anything other than a direct call, and
stay silent otherwise. Decorated definitions keep warning, since a decorator
can store the function.

Fixes #468
Fixes #380

* Tie the B023 exemption to the definition, not to the identifier

Addresses the five review points.

The exemption is now granted only when every reference to the name, in
the scope that holds the loop, is a call that runs in the same iteration
as the definition:

* only a plain `def` qualifies. Calling an `async def` builds a
  coroutine and calling a generator function builds a generator, so the
  body -- and the read of the loop variable -- is deferred.
* a name that is bound anywhere else in the scope is skipped, so a
  reference to an unrelated binding of the same identifier can no longer
  decide the outcome.
* the search covers the loop body only. The `else` suite runs after the
  loop, and a call placed above the `def` invokes the binding the
  previous iteration left behind.
* a reference reached through a nested function, lambda or generator
  expression disqualifies: that body decides when the call happens.
* a definition nothing refers to is reported, because its name is still
  bound after the loop.
* the parent links, name uses and call targets of a scope are built once
  and cached, instead of two walks of the scope per loop.

---------

Co-authored-by: Cooper Lees <me@cooperlees.com>

Author: Eljees

README.rst

@@ -501,6 +501,8 @@ UNRELEASED
 ~~~~~~~~~~
 
 * B019: also flag `async_lru.alru_cache` and check cache decorators on `async def` methods (#488)
+* B023: don't flag a function whose every reference is a direct call inside the loop body:
+  such a function cannot outlive the iteration it was defined in (#468, #380)
 * B020: don't flag `for self.a in self.b`: rebinding an attribute is not rebinding the
   base name, so two different attributes of the same object are two bindings (#248)
 * B018: handle also useless calls such as `isinstance(x, int)` without assigning or using the result

bugbear.py

@@ -76,6 +76,29 @@ class Context(NamedTuple):
     stack: list[ast.AST]
 
 
+def _defines_a_generator(function_node: ast.FunctionDef) -> bool:
+    """Does this function's own body contain a yield? (nested ones do not count)"""
+    stack = list(function_node.body)
+    while stack:
+        node = stack.pop()
+        if isinstance(node, (ast.Yield, ast.YieldFrom)):
+            return True
+        if isinstance(node, FUNCTION_NODES):
+            continue
+        stack.extend(ast.iter_child_nodes(node))
+    return False
+
+
+def _is_rebound(
+    name: str, names_used: dict[str, list[ast.Name]], parents: dict[int, ast.AST]
+) -> bool:
+    """Is `name` bound anywhere in this scope other than by the definition?"""
+    return any(
+        isinstance(reference.ctx, (ast.Store, ast.Del))
+        for reference in names_used.get(name, ())
+    )
+
+
 @attr.s(unsafe_hash=False)
 class BugBearChecker:
     name = "flake8-bugbear"
@@ -428,6 +451,7 @@ class BugBearVisitor(ast.NodeVisitor):
 
     NODE_WINDOW_SIZE = 4
     _b023_seen: set[ast.Name] = attr.ib(factory=set, init=False)
+    _b023_scopes: dict[int, tuple] = attr.ib(factory=dict, init=False)
     _b005_imports: set[str] = attr.ib(factory=set, init=False)
 
     # set to "*" when inside a try/except*, for correctly printing errors
@@ -1025,6 +1049,7 @@ def check_for_b023(  # noqa: C901
         # 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.
+        immediately_called = self._immediately_called_functions(loop_node)
         safe_functions = []
         suspicious_variables = []
         for node in ast.walk(loop_node):
@@ -1058,6 +1083,15 @@ def check_for_b023(  # noqa: C901
                 if isinstance(node.value, FUNCTION_NODES):
                     safe_functions.append(node.value)
 
+            # a function that is only ever *called* in the loop body cannot
+            # outlive the iteration its free variables were assigned in
+            if (
+                isinstance(node, ast.FunctionDef)
+                and not node.decorator_list
+                and node.name in immediately_called
+            ):
+                safe_functions.append(node)
+
             # find unsafe functions
             if isinstance(node, FUNCTION_NODES) and node not in safe_functions:
                 argnames = {
@@ -1101,6 +1135,142 @@ def check_for_b023(  # noqa: C901
             if err.id in reassigned_in_loop:
                 self.add_error("B023", err, err.id)
 
+    def _immediately_called_functions(
+        self,
+        loop_node: (
+            ast.For
+            | ast.AsyncFor
+            | ast.While
+            | ast.GeneratorExp
+            | ast.SetComp
+            | ast.ListComp
+            | ast.DictComp
+        ),
+    ) -> set[str]:
+        """Names of functions defined in the loop that cannot outlive an iteration.
+
+        A function defined in a loop is only subject to the late-binding gotcha
+        B023 warns about if a reference to it survives the iteration it was
+        created in. So a name is reported here -- and thus exempted -- only when
+        every reference to it, anywhere in the enclosing scope, is a direct call
+        that runs in the same iteration as the definition. Being appended to a
+        list, returned, passed as an argument, or called from a nested function
+        all let the function escape, and keep the warning.
+        """
+        body = getattr(loop_node, "body", None)
+        if not isinstance(body, list):
+            return set()
+
+        # Only a plain `def` runs to completion when it is called. Calling an
+        # `async def` builds a coroutine and calling a generator function builds
+        # a generator: both defer the body, so the free variables are read after
+        # the loop has moved on. A decorator may stash the original as well.
+        candidates: dict[str, list[ast.FunctionDef]] = {}
+        for statement in body:
+            if (
+                isinstance(statement, ast.FunctionDef)
+                and not statement.decorator_list
+                and not _defines_a_generator(statement)
+            ):
+                candidates.setdefault(statement.name, []).append(statement)
+        if not candidates:
+            return set()
+
+        parents, names_used, called_names = self._scope_reference_index(loop_node)
+
+        # the `else` suite of a loop runs once the loop is over, so a call there
+        # cannot be the call that keeps the function inside its iteration
+        in_body = set()
+        for statement in body:
+            for node in ast.walk(statement):
+                in_body.add(id(node))
+
+        safe: set[str] = set()
+        for name, definitions in candidates.items():
+            # matching a reference by its identifier alone is only sound while
+            # the name has exactly one binding in this scope
+            if len(definitions) != 1 or _is_rebound(name, names_used, parents):
+                continue
+            definition = definitions[0]
+            # a definition nothing refers to still leaves its name bound after
+            # the loop, so it can be called later with the last iteration's
+            # values -- exactly the bug, so it stays reported
+            references = names_used.get(name, ())
+            if references and all(
+                self._is_call_within_the_iteration(
+                    reference, definition, loop_node, parents, called_names, in_body
+                )
+                for reference in references
+            ):
+                safe.add(name)
+        return safe
+
+    def _scope_reference_index(
+        self, loop_node: ast.AST
+    ) -> tuple[dict[int, ast.AST], dict[str, list[ast.Name]], set[int]]:
+        """Parent links, name uses and call targets for the scope of `loop_node`.
+
+        Built once per scope and cached: walking the scope again for every loop
+        it contains turns a linear traversal into a quadratic one on files with
+        many sequential loops.
+        """
+        scope: ast.AST = loop_node
+        for ancestor in reversed(self.node_stack):
+            if isinstance(ancestor, (ast.Module, ast.ClassDef, *FUNCTION_NODES)):
+                scope = ancestor
+                break
+
+        cached = self._b023_scopes.get(id(scope))
+        if cached is not None:
+            return cached
+
+        parents: dict[int, ast.AST] = {}
+        names_used: dict[str, list[ast.Name]] = {}
+        called_names: set[int] = set()
+        for parent in ast.walk(scope):
+            if isinstance(parent, ast.Call) and isinstance(parent.func, ast.Name):
+                called_names.add(id(parent.func))
+            for child in ast.iter_child_nodes(parent):
+                parents[id(child)] = parent
+            if isinstance(parent, ast.Name):
+                names_used.setdefault(parent.id, []).append(parent)
+
+        index = (parents, names_used, called_names)
+        self._b023_scopes[id(scope)] = index
+        return index
+
+    def _is_call_within_the_iteration(
+        self,
+        reference: ast.Name,
+        definition: ast.FunctionDef,
+        loop_node: ast.AST,
+        parents: dict[int, ast.AST],
+        called_names: set[int],
+        in_body: set[int],
+    ) -> bool:
+        """Is this reference a call that runs in the iteration that defined it?"""
+        if id(reference) not in in_body:
+            return False
+        if not isinstance(reference.ctx, ast.Load):
+            return False
+        if id(reference) not in called_names:
+            return False
+        # a call placed above the `def` invokes the binding the previous
+        # iteration left behind, which is the very bug B023 is about
+        if (reference.lineno, reference.col_offset) < (
+            definition.lineno,
+            definition.col_offset,
+        ):
+            return False
+        # a call reached through another deferred body happens whenever that
+        # body is run, which may be long after the loop has finished
+        node: ast.AST | None = parents.get(id(reference))
+        while node is not None and node is not loop_node:
+            if isinstance(node, (*FUNCTION_NODES, ast.GeneratorExp)):
+                return False
+            node = parents.get(id(node))
+        return True
+
     def check_for_b024_and_b027(self, node: ast.ClassDef) -> None:  # noqa: C901
         """Check for inheritance from abstract classes in abc and lack of
         any methods decorated with abstract*"""

tests/eval_files/b023.py

@@ -169,4 +169,107 @@ def iter_f(names):
             return [lambda: name]  # known false alarm  # B023: 28, "name"
 
         if False:
-            return [lambda: i for i in range(3)]  # error  # B023: 28, "i"
+            return [lambda: i for i in range(3)]  # error  # B023: 28, "i"
+
+# OK because the function is only ever *called* inside the loop body, so it can
+# never outlive the iteration in which its free variables were assigned.
+# https://github.com/PyCQA/flake8-bugbear/issues/468
+for _ in range(10):
+    foo = []
+
+    def immediately_called():
+        foo.append(42)
+
+    immediately_called()
+
+
+# still an error: the function escapes the iteration even though it is also called
+for _ in range(10):
+    bar = []
+
+    def called_and_escapes():
+        bar.append(42)  # B023: 8, "bar"
+
+    called_and_escapes()
+    functions.append(called_and_escapes)
+
+
+# still an error: a decorator can stash the original function somewhere
+for _ in range(10):
+    baz = []
+
+    @some_decorator
+    def decorated():
+        baz.append(42)  # B023: 8, "baz"
+
+    decorated()
+
+
+# still an error: the call happens whenever the *outer* function runs, which may
+# be long after the loop finished
+for _ in range(10):
+    qux = []
+
+    def target():
+        qux.append(42)  # B023: 8, "qux"
+
+    # (`target` itself is not reported: `_get_assigned_names` does not treat a
+    # `def` as an assignment -- pre-existing behaviour, unrelated to this fix)
+    def wrapper():
+        target()
+
+    functions.append(wrapper)
+
+
+# still an error: the function is also called after the loop, so the binding it
+# closes over is whatever the last iteration left behind
+for _ in range(10):
+    quux = []
+
+    def called_after_the_loop():
+        quux.append(42)  # B023: 8, "quux"
+
+    called_after_the_loop()
+
+called_after_the_loop()
+
+# still an error: calling an `async def` only builds a coroutine, so the body --
+# and with it the read of the loop variable -- runs whenever it is awaited
+for _ in range(10):
+    corge = []
+
+    async def awaited_later():
+        corge.append(42)  # B023: 8, "corge"
+
+    awaited_later()
+
+
+# still an error: calling a generator function only builds a generator
+for _ in range(10):
+    grault = []
+
+    def iterated_later():
+        yield grault  # B023: 14, "grault"
+
+    iterated_later()
+
+
+# still an error: a call above the `def` invokes the binding the previous
+# iteration left behind
+for _ in range(10):
+    waldo = []
+
+    called_above_the_def()
+
+    def called_above_the_def():
+        waldo.append(42)  # B023: 8, "waldo"
+
+
+# still an error: the `else` suite runs once the loop is over
+for _ in range(10):
+    garply = []
+
+    def called_in_the_else_suite():
+        garply.append(42)  # B023: 8, "garply"
+else:
+    called_in_the_else_suite()