Repository: iris-hep/func_adl

Issue:
Lambda closure prefers global over local arg name
In the following code:

```python
sys_muon = (calib_tools.query_sys_error(ds, 'MUON_ID__1up')
           .Select(lambda e: e.Muons())
           .Select(lambda l_muons: [m for m in l_muons if m.pt() > 30000.0])
           .Select(lambda l_muons: {
               'pt': [m.pt() / 1000.0 for m in l_muons],
               'eta': [m.eta() for m in l_muons],
               'phi': [m.phi() for m in l_muons],
           })
           .AsAwkwardArray()
           .value())
```

It will pick up `l_muons` if available in local variables or global space rather than the `lambda` argument. Totally violates how we think of python!

Discussion:

Patch:
diff --git a/func_adl/util_ast.py b/func_adl/util_ast.py
index 868029d..7b65cef 100644
--- a/func_adl/util_ast.py
+++ b/func_adl/util_ast.py
@@ -322,14 +322,28 @@ class _rewrite_captured_vars(ast.NodeTransformer):
     def __init__(self, cv: inspect.ClosureVars):
         self._lookup_dict: Dict[str, Any] = dict(cv.nonlocals)
         self._lookup_dict.update(cv.globals)
+        self._ignore_stack = []
 
     def visit_Name(self, node: ast.Name) -> Any:
+        if self.is_arg(node.id):
+            return node
+
         if node.id in self._lookup_dict:
             v = self._lookup_dict[node.id]
             if not callable(v):
                 return as_literal(self._lookup_dict[node.id])
         return node
 
+    def visit_Lambda(self, node: ast.Lambda) -> Any:
+        self._ignore_stack.append([a.arg for a in node.args.args])
+        v = super().generic_visit(node)
+        self._ignore_stack.pop()
+        return v
+
+    def is_arg(self, a_name: str) -> bool:
+        "If the arg is on the stack, then return true"
+        return any([a == a_name for frames in self._ignore_stack for a in frames])
+
 
 def global_getclosurevars(f: Callable) -> inspect.ClosureVars:
     """Grab the closure over all passed function. Add all known global


Test patch:
diff --git a/tests/test_util_ast.py b/tests/test_util_ast.py
index e255a40..5d2284a 100644
--- a/tests/test_util_ast.py
+++ b/tests/test_util_ast.py
@@ -279,9 +279,23 @@ def test_parse_lambda_capture():
     assert ast.dump(r) == ast.dump(r_true)
 
 
+def test_parse_lambda_capture_ignore_local():
+    x = 30  # NOQA type: ignore
+    r = parse_as_ast(lambda x: x > 20)
+    r_true = parse_as_ast(lambda y: y > 20)
+    assert ast.dump(r) == ast.dump(r_true).replace("'y'", "'x'")
+
+
 g_cut_value = 30
 
 
+def test_parse_lambda_capture_ignore_global():
+    x = 30  # NOQA type: ignore
+    r = parse_as_ast(lambda g_cut_value: g_cut_value > 20)
+    r_true = parse_as_ast(lambda y: y > 20)
+    assert ast.dump(r) == ast.dump(r_true).replace("'y'", "'g_cut_value'")
+
+
 def test_parse_lambda_capture_nested_global():
     r = parse_as_ast(lambda x: (lambda y: y > g_cut_value)(x))
     r_true = parse_as_ast(lambda x: (lambda y: y > 30)(x))