Repository: iris-hep/func_adl

Issue:
Variables in lambdas with identical names in a query
Using `func-adl` version `3.2.2`, the following query works:
```python
ds.Where(lambda e: e.met_e > 25).Select(lambda event: event.jet_pt)
```
but the same query does no longer work when the variables in both lambdas have the same name:
```python
ds.Where(lambda e: e.met_e > 25).Select(lambda e: e.jet_pt)
```
raises

> ValueError: Found multiple calls to on same line - split the calls across lines or change lambda argument names so they are different.

Splitting across lines does not seem to help, e.g.
```python
ds.Where(lambda e: e.met_e > 25).\
Select(lambda e: e.jet_pt)
```
similarly crashes.

---

Reproducer:
```python
from func_adl_uproot import UprootDataset

dataset_opendata = "http://xrootd-local.unl.edu:1094//store/user/AGC/datasets/RunIIFall15MiniAODv2/"\
                   "TT_TuneCUETP8M1_13TeV-powheg-pythia8/MINIAODSIM//PU25nsData2015v1_76X_mcRun2_"\
                   "asymptotic_v12_ext3-v1/00000/00DF0A73-17C2-E511-B086-E41D2D08DE30.root"
ds = UprootDataset(dataset_opendata, "events")
jet_pt_query = ds.Where(lambda e: e.met_e > 25).Select(lambda ev: ev.jet_pt)
print(jet_pt_query.value())
```
I used `func-adl-uproot` version `1.10.0`. Full traceback:
```pytb
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [25], in <cell line: 6>()
      3 dataset_opendata = "http://xrootd-local.unl.edu:1094//store/user/AGC/datasets/RunIIFall15MiniAODv2/TT_TuneCUETP8M1_13TeV-powheg-pythia8/"\
      4                    "MINIAODSIM//PU25nsData2015v1_76X_mcRun2_asymptotic_v12_ext3-v1/00000/00DF0A73-17C2-E511-B086-E41D2D08DE30.root"
      5 ds = UprootDataset(dataset_opendata, "events")
----> 6 jet_pt_query = ds.Where(lambda e: e.met_e > 25).Select(lambda e: e.jet_pt)
      7 print(jet_pt_query.value())

File /opt/conda/lib/python3.8/site-packages/func_adl/object_stream.py:163, in ObjectStream.Where(self, filter)
    145 r"""
    146 Filter the object stream, allowing only items for which `filter` evaluates as true through.
    147 
   (...)
    158       contains a lambda definition, or a python `ast` of type `ast.Lambda`.
    159 """
    160 from func_adl.type_based_replacement import remap_from_lambda
    162 n_stream, n_ast, rtn_type = remap_from_lambda(
--> 163     self, _local_simplification(parse_as_ast(filter))
    164 )
    165 if rtn_type != bool:
    166     raise ValueError(f"The Where filter must return a boolean (not {rtn_type})")

File /opt/conda/lib/python3.8/site-packages/func_adl/util_ast.py:683, in parse_as_ast(ast_source, caller_name)
    664 r"""Return an AST for a lambda function from several sources.
    665 
    666 We are handed one of several things:
   (...)
    680     An ast starting from the Lambda AST node.
    681 """
    682 if callable(ast_source):
--> 683     src_ast = _parse_source_for_lambda(ast_source, caller_name)
    684     if not src_ast:
    685         # This most often happens in a notebook when the lambda is defined in a funny place
    686         # and can't be recovered.
    687         raise ValueError(f"Unable to recover source for function {ast_source}.")

File /opt/conda/lib/python3.8/site-packages/func_adl/util_ast.py:649, in _parse_source_for_lambda(ast_source, caller_name)
    644         raise ValueError(
    645             f"Internal Error - Found no lambda in source with the arguments {caller_arg_list}"
    646         )
    648     if len(good_lambdas) > 1:
--> 649         raise ValueError(
    650             "Found multiple calls to on same line"
    651             + ("" if caller_name is None else f" for {caller_name}")
    652             + " - split the calls across "
    653             "lines or change lambda argument names so they are different."
    654         )
    656     lda = good_lambdas[0]
    658 return lda

ValueError: Found multiple calls to on same line - split the calls across lines or change lambda argument names so they are different.
```

Comments:

Patch:
diff --git a/func_adl/object_stream.py b/func_adl/object_stream.py
index f25cec0..09a4410 100644
--- a/func_adl/object_stream.py
+++ b/func_adl/object_stream.py
@@ -110,7 +110,7 @@ class ObjectStream(Generic[T]):
         from func_adl.type_based_replacement import remap_from_lambda
 
         n_stream, n_ast, rtn_type = remap_from_lambda(
-            self, _local_simplification(parse_as_ast(func))
+            self, _local_simplification(parse_as_ast(func, "SelectMany"))
         )
         return ObjectStream[S](
             function_call("SelectMany", [n_stream.query_ast, cast(ast.AST, n_ast)]),
@@ -136,7 +136,9 @@ class ObjectStream(Generic[T]):
         """
         from func_adl.type_based_replacement import remap_from_lambda
 
-        n_stream, n_ast, rtn_type = remap_from_lambda(self, _local_simplification(parse_as_ast(f)))
+        n_stream, n_ast, rtn_type = remap_from_lambda(
+            self, _local_simplification(parse_as_ast(f, "Select"))
+        )
         return ObjectStream[S](
             function_call("Select", [n_stream.query_ast, cast(ast.AST, n_ast)]), rtn_type
         )
@@ -160,7 +162,7 @@ class ObjectStream(Generic[T]):
         from func_adl.type_based_replacement import remap_from_lambda
 
         n_stream, n_ast, rtn_type = remap_from_lambda(
-            self, _local_simplification(parse_as_ast(filter))
+            self, _local_simplification(parse_as_ast(filter, "Where"))
         )
         if rtn_type != bool:
             raise ValueError(f"The Where filter must return a boolean (not {rtn_type})")
diff --git a/func_adl/util_ast.py b/func_adl/util_ast.py
index f4afa12..5c5f8a4 100644
--- a/func_adl/util_ast.py
+++ b/func_adl/util_ast.py
@@ -647,10 +647,14 @@ def _parse_source_for_lambda(
 
         if len(good_lambdas) > 1:
             raise ValueError(
-                "Found multiple calls to on same line"
+                "Found multiple calls on same line"
                 + ("" if caller_name is None else f" for {caller_name}")
                 + " - split the calls across "
-                "lines or change lambda argument names so they are different."
+                """lines or change lambda argument names so they are different. For example change:
+                    df.Select(lambda x: x + 1).Select(lambda x: x + 2)
+                    to:
+                    df.Select(lambda x: x + 1).Select(lambda y: y + 2)
+                """
             )
 
         lda = good_lambdas[0]


Test patch:
diff --git a/tests/test_object_stream.py b/tests/test_object_stream.py
index f418451..a1ed255 100644
--- a/tests/test_object_stream.py
+++ b/tests/test_object_stream.py
@@ -75,6 +75,12 @@ def test_simple_query():
     assert isinstance(r, ast.AST)
 
 
+def test_simple_query_one_line():
+    """Make sure we parse 2 functions on one line correctly"""
+    r = my_event().Select(lambda e: e.met).Where(lambda e: e > 10).value()
+    assert isinstance(r, ast.AST)
+
+
 def test_two_simple_query():
     r1 = (
         my_event()
diff --git a/tests/test_util_ast.py b/tests/test_util_ast.py
index 9e8d86c..e4a8ee1 100644
--- a/tests/test_util_ast.py
+++ b/tests/test_util_ast.py
@@ -682,6 +682,38 @@ def test_parse_multiline_lambda_blank_lines_no_infinite_loop():
     assert "uncalibrated_collection" in ast.dump(found[0])
 
 
+def test_parse_select_where():
+    "Common lambas with different parent functions on one line - found in wild"
+
+    found = []
+
+    class my_obj:
+        def Where(self, x: Callable):
+            found.append(parse_as_ast(x, "Where"))
+            return self
+
+        def Select(self, x: Callable):
+            found.append(parse_as_ast(x, "Select"))
+            return self
+
+        def AsAwkwardArray(self, stuff: str):
+            return self
+
+        def value(self):
+            return self
+
+    jets_pflow_name = "hi"
+    ds_dijet = my_obj()
+
+    # fmt: off
+    jets_pflow = (
+        ds_dijet.Select(lambda e: e.met).Where(lambda e: e > 100)
+    )
+    # fmt: on
+    assert jets_pflow is not None  # Just to keep flake8 happy without adding a noqa above.
+    assert "met" in ast.dump(found[0])
+
+
 def test_parse_multiline_lambda_ok_with_one_as_arg():
     "Make sure we can properly parse a multi-line lambda - but now with argument"