Use Function Names as Clues Parsing Lambda's (#116)

* Make sure we use the fact we are after a `Where` or `Select` when parsing multiple `lambda`'s on a line
* Improve the error message when we can't tell the difference between mtuple lambda's.

Fixes #115

Author: gordonwatts

func_adl/object_stream.py

@@ -110,7 +110,7 @@ def SelectMany(
         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 @@ def Select(self, f: Union[str, ast.Lambda, Callable[[T], S]]) -> ObjectStream[S]
         """
         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 @@ def Where(self, filter: Union[str, ast.Lambda, Callable[[T], bool]]) -> ObjectSt
         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})")

func_adl/util_ast.py

@@ -647,10 +647,14 @@ def lambda_arg_list(lda: ast.Lambda) -> List[str]:
 
         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]

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()

tests/test_util_ast.py

@@ -682,6 +682,38 @@ def value(self):
     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"