Quoted strings for lambda removed. (#204)

Unit tests should not use quoted strings for func_adl. No advantage internally but does make for better testing of the `func_adl` `labmda` extractor. Fixes #173
iris-hep · Nov 13, 2022 · ebee8e6 · ebee8e6
1 parent dd380cf
commit ebee8e6
Show file tree

Hide file tree

Showing 14 changed files with 679 additions and 371 deletions.
diff --git a/scripts/generate_ast_text.py b/scripts/generate_ast_text.py
@@ -16,32 +16,53 @@ async def execute_result_async(self, a: ast.AST, title: str) -> Any:
 
 
 def as_ast_lang_query_0():
-    'Extract jet pt of all jets'
-    r = ast_translator_ds() \
-        .SelectMany("lambda e: e.Jets('')") \
-        .Select("lambda j: j.pt()") \
-        .AsROOTTTree("junk.root", 'analysis', ['jet_pt']) \
+    "Extract jet pt of all jets"
+    r = (
+        ast_translator_ds()
+        .SelectMany(lambda e: e.Jets(""))
+        .Select(lambda j: j.pt())
+        .AsROOTTTree("junk.root", "analysis", ["jet_pt"])
         .value()
+    )
     return r
 
 
 def as_ast_lang_query_1():
-    'Generate a real query that works on our 10 event root file'
-    r = ast_translator_ds() \
-        .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets")') \
-        .Select('lambda j: j.pt()/1000.0') \
-        .AsROOTTTree("junk.root", "analysis", ['JetPt']) \
+    "Generate a real query that works on our 10 event root file"
+    r = (
+        ast_translator_ds()
+        .SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets"))
+        .Select(lambda j: j.pt() / 1000.0)
+        .AsROOTTTree("junk.root", "analysis", ["JetPt"])
         .value()
+    )
     return r
 
 
 def as_ast_lang_query_2():
-    'Marc needed this to run some performance tests.'
-    r = ast_translator_ds() \
-        .Select('lambda e: (e.Electrons("Electrons"), e.Muons("Muons"))') \
-        .Select('lambda e: (e[0].Select(lambda ele: ele.e()), e[0].Select(lambda ele: ele.pt()), e[0].Select(lambda ele: ele.phi()), e[0].Select(lambda ele: ele.eta()), e[1].Select(lambda mu: mu.e()), e[1].Select(lambda mu: mu.pt()), e[1].Select(lambda mu: mu.phi()), e[1].Select(lambda mu: mu.eta()))') \
-        .AsROOTTTree('dude.root', 'forkme', ['e_E', 'e_pt', 'e_phi', 'e_eta', 'mu_E', 'mu_pt', 'mu_phi', 'mu_eta']) \
+    "Marc needed this to run some performance tests."
+    r = (
+        ast_translator_ds()
+        .Select(lambda e: (e.Electrons("Electrons"), e.Muons("Muons")))
+        .Select(
+            lambda e: (
+                e[0].Select(lambda ele: ele.e()),
+                e[0].Select(lambda ele: ele.pt()),
+                e[0].Select(lambda ele: ele.phi()),
+                e[0].Select(lambda ele: ele.eta()),
+                e[1].Select(lambda mu: mu.e()),
+                e[1].Select(lambda mu: mu.pt()),
+                e[1].Select(lambda mu: mu.phi()),
+                e[1].Select(lambda mu: mu.eta()),
+            )
+        )
+        .AsROOTTTree(
+            "dude.root",
+            "forkme",
+            ["e_E", "e_pt", "e_phi", "e_eta", "mu_E", "mu_pt", "mu_phi", "mu_eta"],
+        )
         .value()
+    )
     return r
 
 

diff --git a/tests/atlas/r22_xaod/test_integrated_query_r22.py b/tests/atlas/r22_xaod/test_integrated_query_r22.py
@@ -36,8 +36,8 @@ def event_loop():
 def test_flatten_array():
     # A very simple flattening of arrays
     training_df = as_pandas(
-        f_single.SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets")').Select(
-            "lambda j: j.pt()/1000.0"
+        f_single.SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets")).Select(
+            lambda j: j.pt() / 1000.0
         )
     )
     assert abs(training_df.iloc[0]["col1"] - 21.733) < 0.001  # type: ignore ()

diff --git a/tests/atlas/xaod/test_Jets.py b/tests/atlas/xaod/test_Jets.py
@@ -6,9 +6,15 @@
 
 def test_get_attribute_float():
     # The following statement should be a straight sequence, not an array.
-    r = atlas_xaod_dataset() \
-        .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.getAttributeFloat("emf"))') \
+    r = (
+        atlas_xaod_dataset()
+        .SelectMany(
+            lambda e: e.Jets("AntiKt4EMTopoJets").Select(
+                lambda j: j.getAttributeFloat("emf")
+            )
+        )
         .value()
+    )
     # Check to see if there mention of push_back anywhere.
     lines = get_lines_of_code(r)
     print_lines(lines)
@@ -18,19 +24,23 @@ def test_get_attribute_float():
 
 def test_get_attribute_float_wrong_args():
     with pytest.raises(Exception) as e:
-        atlas_xaod_dataset() \
-            .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.getAttributeFloat())') \
-            .value()
+        atlas_xaod_dataset().SelectMany(
+            lambda e: e.Jets("AntiKt4EMTopoJets").Select(
+                lambda j: j.getAttributeFloat()
+            )
+        ).value()
 
-    assert 'getAttribute' in str(e.value)
+    assert "getAttribute" in str(e.value)
 
 
 def test_get_attribute_vector_float():
     # The following statement should be a straight sequence, not an array.
-    r = atlas_xaod_dataset() \
-        .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets")') \
-        .SelectMany('lambda j: j.getAttributeVectorFloat("emf")') \
+    r = (
+        atlas_xaod_dataset()
+        .SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets"))
+        .SelectMany(lambda j: j.getAttributeVectorFloat("emf"))
         .value()
+    )
     # Check to see if there mention of push_back anywhere.
     lines = get_lines_of_code(r)
     print_lines(lines)
@@ -40,8 +50,10 @@ def test_get_attribute_vector_float():
 
 def test_get_attribute():
     with pytest.raises(Exception) as e:
-        atlas_xaod_dataset() \
-            .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.getAttribute("emf"))') \
-            .value()
+        atlas_xaod_dataset().SelectMany(
+            lambda e: e.Jets("AntiKt4EMTopoJets").Select(
+                lambda j: j.getAttribute("emf")
+            )
+        ).value()
 
     assert "getAttributeFloat" in str(e.value)
diff --git a/tests/atlas/xaod/test_atlas_xaod_exectuor.py b/tests/atlas/xaod/test_atlas_xaod_exectuor.py
@@ -6,7 +6,7 @@
 
 
 def test_ctor():
-    'Make sure that the ctor works'
+    "Make sure that the ctor works"
     atlas_xaod_executor()
 
 
@@ -16,12 +16,10 @@ async def execute_result_async(self, a: ast.AST, title: str):
 
 
 def test_xaod_executor(tmp_path):
-    'Write out C++ files for a simple query'
+    "Write out C++ files for a simple query"
 
     # Get the ast to play with
-    a = query_as_ast() \
-        .Select('lambda e: e.EventInfo("EventInfo").runNumber()') \
-        .value()
+    a = query_as_ast().Select(lambda e: e.EventInfo("EventInfo").runNumber()).value()
 
     exe = atlas_xaod_executor()
     f_spec = exe.write_cpp_files(exe.apply_ast_transformations(a), tmp_path)
@@ -30,46 +28,42 @@ def test_xaod_executor(tmp_path):
 
 
 def test_xaod_library_there(tmp_path):
-    'Make sure a required library is in the link list'
+    "Make sure a required library is in the link list"
     # Get the ast to play with
-    a = query_as_ast() \
-        .Select('lambda e: e.EventInfo("EventInfo").runNumber()') \
-        .value()
+    a = query_as_ast().Select(lambda e: e.EventInfo("EventInfo").runNumber()).value()
 
     exe = atlas_xaod_executor()
     exe.write_cpp_files(exe.apply_ast_transformations(a), tmp_path)
 
-    make_list = tmp_path / 'package_CMakeLists.txt'
+    make_list = tmp_path / "package_CMakeLists.txt"
     assert make_list.exists()
-    assert 'xAODEventInfo' in make_list.read_text()
+    assert "xAODEventInfo" in make_list.read_text()
 
 
 def test_eventinfo_handle_code(tmp_path):
-    'Make sure a required library is in the link list'
+    "Make sure a required library is in the link list"
     # Get the ast to play with
-    a = query_as_ast() \
-        .Select('lambda e: e.EventInfo("EventInfo").runNumber()') \
-        .value()
+    a = query_as_ast().Select(lambda e: e.EventInfo("EventInfo").runNumber()).value()
 
     exe = atlas_xaod_executor()
     exe.write_cpp_files(exe.apply_ast_transformations(a), tmp_path)
 
-    query = tmp_path / 'query.cxx'
-    assert 'const xAOD::EventInfo *' in query.read_text()
+    query = tmp_path / "query.cxx"
+    assert "const xAOD::EventInfo *" in query.read_text()
 
 
 def test_find_exception():
-    'Make sure _find exception is well formed'
+    "Make sure _find exception is well formed"
     from func_adl_xAOD.common.executor import _find
 
     with pytest.raises(RuntimeError) as e:
-        _find('fork-it-over.txt')
+        _find("fork-it-over.txt")
 
-    assert 'find file' in str(e.value)
+    assert "find file" in str(e.value)
 
 
 def test_bad_ast_no_call(tmp_path):
-    'Pass a really bogus ast to the executor'
+    "Pass a really bogus ast to the executor"
     # Get the ast to play with
     q = query_as_ast()
     a = ast.UnaryOp(op=ast.USub(), operand=q.query_ast)
@@ -78,63 +72,73 @@ def test_bad_ast_no_call(tmp_path):
     with pytest.raises(ValueError) as e:
         exe.write_cpp_files(exe.apply_ast_transformations(a), tmp_path)
 
-    assert 'func_adl ast' in str(e.value)
+    assert "func_adl ast" in str(e.value)
 
 
 def test_bad_ast_no_call_to_name(tmp_path):
-    'Pass a really bogus ast to the executor'
+    "Pass a really bogus ast to the executor"
     # Get the ast to play with
     q = query_as_ast()
-    a = ast.Call(func=ast.Attribute(value=ast.Constant(10), attr='fork'), args=[q.query_ast])
+    a = ast.Call(
+        func=ast.Attribute(value=ast.Constant(10), attr="fork"), args=[q.query_ast]
+    )
 
     exe = atlas_xaod_executor()
     with pytest.raises(ValueError) as e:
         exe.write_cpp_files(exe.apply_ast_transformations(a), tmp_path)
 
-    assert 'func_adl ast' in str(e.value)
+    assert "func_adl ast" in str(e.value)
 
 
 def test_md_job_options(tmp_path):
-    'Make sure our job options script appears in the right place'
+    "Make sure our job options script appears in the right place"
 
     # Get the ast to play with
-    a = (query_as_ast()
-         .MetaData({
-                   'metadata_type': 'add_job_script',
-                   'name': 'Vertex',
-                   'script': [
-                       '# this is a fork tester',
-                   ]
-                   })
-         .Select('lambda e: e.EventInfo("EventInfo").runNumber()')
-         .value())
+    a = (
+        query_as_ast()
+        .MetaData(
+            {
+                "metadata_type": "add_job_script",
+                "name": "Vertex",
+                "script": [
+                    "# this is a fork tester",
+                ],
+            }
+        )
+        .Select(lambda e: e.EventInfo("EventInfo").runNumber())
+        .value()
+    )
 
     exe = atlas_xaod_executor()
     exe.write_cpp_files(exe.apply_ast_transformations(a), tmp_path)
 
-    with open(tmp_path / 'ATestRun_eljob.py', 'r') as f:
+    with open(tmp_path / "ATestRun_eljob.py", "r") as f:
         lines = [ln.strip() for ln in f.readlines()]
-        assert '# this is a fork tester' in lines
+        assert "# this is a fork tester" in lines
 
 
 def test_md_replaced_collection(tmp_path):
-    'When we replace a collection, make sure all goes right'
+    "When we replace a collection, make sure all goes right"
     # Get the ast to play with
-    a = query_as_ast() \
-        .MetaData({
-            'metadata_type': 'add_atlas_event_collection_info',
-            'name': 'EventInfo',
-            'include_files': ['xAODEventInfo/versions/EventInfo_v1.h'],
-            'container_type': 'xAOD::EventInfo_v1',
-            'contains_collection': False,
-            'link_libraries': ['xAODEventInfo'],
-        }) \
-        .Select('lambda e: e.EventInfo("EventInfo").runNumber()') \
+    a = (
+        query_as_ast()
+        .MetaData(
+            {
+                "metadata_type": "add_atlas_event_collection_info",
+                "name": "EventInfo",
+                "include_files": ["xAODEventInfo/versions/EventInfo_v1.h"],
+                "container_type": "xAOD::EventInfo_v1",
+                "contains_collection": False,
+                "link_libraries": ["xAODEventInfo"],
+            }
+        )
+        .Select(lambda e: e.EventInfo("EventInfo").runNumber())
         .value()
+    )
 
     exe = atlas_xaod_executor()
     exe.write_cpp_files(exe.apply_ast_transformations(a), tmp_path)
 
-    query = tmp_path / 'query.cxx'
+    query = tmp_path / "query.cxx"
     assert query.exists()
-    assert 'const xAOD::EventInfo_v1 *' in query.read_text()
+    assert "const xAOD::EventInfo_v1 *" in query.read_text()
diff --git a/tests/atlas/xaod/test_cpp_script_control.py b/tests/atlas/xaod/test_cpp_script_control.py
@@ -45,8 +45,8 @@ def generate_test_jet_fetch(cache_dir: Path):
     """
     return (
         hash_event_dataset(cache_dir)
-        .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets")')
-        .Select("lambda j: j.pt()/1000.0")
+        .SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets"))
+        .Select(lambda j: j.pt() / 1000.0)
         .value()
     )
 
@@ -57,8 +57,8 @@ def generate_test_jet_fetch_bad(cache_dir: Path):
     """
     return (
         hash_event_dataset(cache_dir)
-        .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets")')
-        .Select("lambda j: j.ptt()/1000.0")
+        .SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets"))
+        .Select(lambda j: j.ptt() / 1000.0)
         .value()
     )