fix[next]: collect closure vars via the compiler's scope analysis

src/gt4py/next/ffront/source_utils.py

@@ -8,6 +8,7 @@
 
 from __future__ import annotations
 
+import builtins
 import functools
 import inspect
 import pathlib
@@ -21,11 +22,48 @@
 MISSING_FILENAME = "<string>"
 
 
+def _global_names_from_source(source: str) -> set[str]:
+    """
+    Names referenced as globals in the function defined by `source`, including its nested scopes.
+
+    The compiler's own scope analysis decides what is a global, so locals shadowing a
+    global name and comprehension targets are never reported.
+    """
+
+    def walk(table: symtable.SymbolTable) -> Iterator[str]:
+        if isinstance(table, symtable.Function):
+            yield from table.get_globals()
+        for child in table.get_children():
+            yield from walk(child)
+
+    # Analyzed as if under PEP 563 so that annotations contribute no names: Python evaluates
+    # parameter and return annotations in the enclosing scope and never evaluates local
+    # variable annotations, but only Python 3.14's symtable stops reporting the latter.
+    source = "from __future__ import annotations\n" + source
+    return set(walk(symtable.symtable(source, MISSING_FILENAME, "exec")))
+
+
 def get_closure_vars_from_function(function: Callable) -> dict[str, Any]:
-    (nonlocals, globals, builtins, _unbound) = inspect.getclosurevars(function)  # noqa: A001 [builtin-variable-shadowing]
+    # `inspect.getclosurevars` only sees the names of the function's own code object, which
+    # misses names referenced only inside a nested scope such as a generator expression.
+    # Free variables are unaffected (they are cells of the function itself), so only the
+    # global names are taken from the source instead.
+    nonlocals = inspect.getclosurevars(function).nonlocals
+    global_ns = function.__globals__
+    builtin_ns = global_ns.get("__builtins__", builtins.__dict__)
+    if inspect.ismodule(builtin_ns):
+        builtin_ns = builtin_ns.__dict__
+
+    source = make_source_definition_from_function(function).source
+    closure_vars: dict[str, Any] = {}
+    for name in _global_names_from_source(source):
+        if name in global_ns:
+            closure_vars[name] = global_ns[name]
+        elif name in builtin_ns:
+            closure_vars[name] = builtin_ns[name]
 
     # nonlocals override globals, sorted for deterministic results
-    return dict(sorted({**builtins, **globals, **nonlocals}.items()))
+    return dict(sorted({**closure_vars, **nonlocals}.items()))
 
 
 def make_source_definition_from_function(func: Callable) -> SourceDefinition:

tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py

@@ -10,7 +10,7 @@
 import pytest
 
 import gt4py.next as gtx
-from gt4py.next import broadcast, errors, float64, int32, neighbor_sum, utils as gt_utils
+from gt4py.next import broadcast, common, errors, float64, int32, neighbor_sum, utils as gt_utils
 
 from next_tests.integration_tests import cases
 from next_tests.integration_tests.cases import (
@@ -199,6 +199,63 @@ def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IFie
     )
 
 
+# Defined at module level, unlike `inner` in 'test_tuple_comprehension_other_fo' above.
+# That difference alone decides whether the comprehension below can see it.
+@gtx.field_operator
+def _module_level_inner(tracer: cases.IField, factor: int32) -> cases.IField:
+    return tracer * factor
+
+
+@pytest.mark.uses_tuple_args
+def test_tuple_comprehension_module_level_fo(cartesian_case):
+    # Identical to 'test_tuple_comprehension_other_fo', except that the called operator is
+    # a module-level global rather than a local of the test function.
+    @gtx.field_operator
+    def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]:
+        return tuple(_module_level_inner(tracer, factor) for tracer in tracers)
+
+    cases.verify_with_default_data(
+        cartesian_case,
+        testee,
+        ref=lambda t, f: tuple(el * f for el in t),
+    )
+
+
+@pytest.mark.uses_tuple_args
+@pytest.mark.uses_unstructured_shift
+def test_tuple_comprehension_module_level_builtin(unstructured_case):
+    # Same cause, but for an imported builtin -- the shape every real stencil has.
+    @gtx.field_operator
+    def testee(tracers: tuple[cases.EField, ...]) -> tuple[cases.VField, ...]:
+        return tuple(neighbor_sum(tracer(V2E), axis=V2EDim) for tracer in tracers)
+
+    v2e = unstructured_case.offset_provider["V2E"].asnumpy()
+    valid = v2e != common._DEFAULT_SKIP_VALUE
+    cases.verify_with_default_data(
+        unstructured_case,
+        testee,
+        ref=lambda t: tuple(np.sum(el[v2e], axis=1, where=valid) for el in t),
+    )
+
+
+@pytest.mark.uses_tuple_args
+def test_tuple_comprehension_module_level_fo_used_outside_too(cartesian_case):
+    # The workaround, and why this is easy to miss: the comprehension body is unchanged;
+    # only the unrelated statement above it makes '_module_level_inner' collectable.
+    @gtx.field_operator
+    def testee(
+        tracers: tuple[cases.IField, ...], a: cases.IField, factor: int32
+    ) -> tuple[cases.IField, ...]:
+        unrelated = _module_level_inner(a, factor)
+        return tuple(_module_level_inner(tracer, factor) for tracer in tracers)
+
+    cases.verify_with_default_data(
+        cartesian_case,
+        testee,
+        ref=lambda t, a, f: tuple(el * f for el in t),
+    )
+
+
 @pytest.mark.uses_tuple_args
 def test_nested_tuple_comprehension(cartesian_case):
     @gtx.field_operator

tests/next_tests/unit_tests/ffront_tests/test_tuple_comprehension_closure_vars.py

@@ -0,0 +1,99 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2024, ETH Zurich
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Closure variable resolution inside tuple comprehension bodies.
+
+A generator expression compiles to its own code object, so names referenced only
+inside it appear in that object's `co_names` and never in the enclosing function's.
+Collecting closure variables from the enclosing function alone therefore misses
+them; `get_closure_vars_from_function` takes the global names from the compiler's
+symbol table of the source instead, which covers nested scopes.
+
+Free variables need no such treatment: the enclosing code object carries a cell
+for them, which is why comprehensions over locally defined helpers always worked.
+"""
+
+import pytest
+
+import gt4py.next as gtx
+from gt4py.next import Dims, Dimension, float64, neighbor_sum
+from gt4py.next.ffront.func_to_foast import FieldOperatorParser
+from gt4py.next.ffront.source_utils import get_closure_vars_from_function
+
+
+Cell = Dimension("Cell")
+Edge = Dimension("Edge")
+C2EDim = Dimension("C2E", kind=gtx.DimensionKind.LOCAL)
+C2E = gtx.FieldOffset("C2E", source=Edge, target=(Cell, C2EDim))
+
+CField = gtx.Field[Dims[Cell], float64]
+EField = gtx.Field[Dims[Edge], float64]
+
+
+@gtx.field_operator
+def scale(f: CField, factor: float64) -> CField:
+    return f * factor
+
+
+def _builtin_and_offset(tracers: tuple[EField, ...]) -> tuple[CField, ...]:
+    return tuple(neighbor_sum(t(C2E), axis=C2EDim) for t in tracers)
+
+
+def _module_level_operator(tracers: tuple[CField, ...], factor: float64) -> tuple[CField, ...]:
+    return tuple(scale(t, factor) for t in tracers)
+
+
+@pytest.mark.parametrize(
+    "func", [_builtin_and_offset, _module_level_operator], ids=["builtin_and_offset", "operator"]
+)
+def test_module_level_name_used_only_in_comprehension(func):
+    """Module-level names are reachable from inside a comprehension body."""
+    parsed = FieldOperatorParser.apply_to_function(func)
+    assert parsed.type is not None
+
+
+def test_names_are_collected_from_the_nested_code_object():
+    """The names live in the generator expression's code object, not the function's."""
+    assert "neighbor_sum" not in _builtin_and_offset.__code__.co_names
+    nested = [c for c in _builtin_and_offset.__code__.co_consts if hasattr(c, "co_names")]
+    assert any("neighbor_sum" in c.co_names for c in nested)
+
+    collected = get_closure_vars_from_function(_builtin_and_offset)
+    assert {"neighbor_sum", "C2E", "C2EDim"} <= set(collected)
+
+
+def test_comprehension_target_is_not_collected():
+    """The loop target is a local of the nested code object, not a global reference."""
+    assert "t" not in get_closure_vars_from_function(_module_level_operator)
+
+
+def test_free_variables_still_resolve():
+    """The path that already worked, kept as a guard."""
+
+    @gtx.field_operator
+    def local_scale(f: CField, factor: float64) -> CField:
+        return f * factor
+
+    def uses_freevar(tracers: tuple[CField, ...], factor: float64) -> tuple[CField, ...]:
+        return tuple(local_scale(t, factor) for t in tracers)
+
+    parsed = FieldOperatorParser.apply_to_function(uses_freevar)
+    assert parsed.type is not None
+
+
+def test_local_name_shadowing_a_global_is_not_collected_as_global():
+    """A comprehension referencing an enclosing local must bind the local, not the global."""
+
+    def shadows(tracers: tuple[CField, ...], factor: float64) -> tuple[CField, ...]:
+        scale = local_helper  # noqa: F841  shadows the module-level 'scale'
+        return tuple(scale(t, factor) for t in tracers)
+
+    def local_helper(t, factor):
+        return t
+
+    assert "scale" not in get_closure_vars_from_function(shadows)