From ebee8e6ff605171762241143334c8d0c2c168f73 Mon Sep 17 00:00:00 2001 From: Gordon Watts Date: Sun, 13 Nov 2022 11:39:54 -0500 Subject: [PATCH] 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 --- scripts/generate_ast_text.py | 51 +++- .../r22_xaod/test_integrated_query_r22.py | 4 +- tests/atlas/xaod/test_Jets.py | 36 ++- tests/atlas/xaod/test_atlas_xaod_exectuor.py | 106 +++---- tests/atlas/xaod/test_cpp_script_control.py | 8 +- tests/atlas/xaod/test_first_last.py | 68 +++-- tests/atlas/xaod/test_integrated_query.py | 4 +- tests/atlas/xaod/test_simple_type_info.py | 41 ++- tests/atlas/xaod/test_xaod_aggragate.py | 173 +++++++++--- tests/atlas/xaod/test_xaod_executor.py | 263 +++++++++++++----- tests/cms/aod/test_integrated_query.py | 97 ++++--- tests/cms/aod/test_local_dataset.py | 42 +-- tests/cms/miniaod/test_integrated_query.py | 115 ++++---- tests/cms/miniaod/test_local_dataset.py | 42 +-- 14 files changed, 679 insertions(+), 371 deletions(-) diff --git a/scripts/generate_ast_text.py b/scripts/generate_ast_text.py index dafdffef..677bf8d5 100644 --- 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 index 9bbb5f7c..bc0d79c8 100644 --- 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 index 607c2b2d..fcdb4c58 100644 --- 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 index ca31ddaa..d431694b 100644 --- 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 index 26da315d..7458e5c5 100644 --- 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() ) diff --git a/tests/atlas/xaod/test_first_last.py b/tests/atlas/xaod/test_first_last.py index 57fc3c6e..5da87b6b 100644 --- a/tests/atlas/xaod/test_first_last.py +++ b/tests/atlas/xaod/test_first_last.py @@ -6,15 +6,21 @@ def test_first_jet_in_event(): - atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("bogus").Select(lambda j: j.pt()).First()') \ - .value() + atlas_xaod_dataset().Select( + lambda e: e.Jets("bogus").Select(lambda j: j.pt()).First() + ).value() def test_first_after_selectmany(): - r = atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("jets").SelectMany(lambda j: e.Tracks("InnerTracks")).First()') \ + r = ( + atlas_xaod_dataset() + .Select( + lambda e: e.Jets("jets") + .SelectMany(lambda j: e.Tracks("InnerTracks")) + .First() + ) .value() + ) lines = get_lines_of_code(r) print_lines(lines) @@ -23,28 +29,34 @@ def test_first_after_where(): # Part of testing that First puts its outer settings in the right place. # This also tests First on a collection of objects that hasn't been pulled a part # in a select. - atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("AntiKt4EMTopoJets").Where(lambda j: j.pt() > 10).First().pt()') \ - .value() + atlas_xaod_dataset().Select( + lambda e: e.Jets("AntiKt4EMTopoJets").Where(lambda j: j.pt() > 10).First().pt() + ).value() def test_first_object_in_each_event(): # Part of testing that First puts its outer settings in the right place. # This also tests First on a collection of objects that hasn't been pulled a part # in a select. - atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("AntiKt4EMTopoJets").First().pt()/1000.0') \ - .value() + atlas_xaod_dataset().Select( + lambda e: e.Jets("AntiKt4EMTopoJets").First().pt() / 1000.0 + ).value() def test_First_Of_Select_is_not_array(): # The following statement should be a straight sequence, not an array. - r = atlas_xaod_dataset() \ - .Select(lambda e: - { - 'FirstJetPt': e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt() / 1000.0).Where(lambda jpt: jpt > 10.0).First() - }) \ + r = ( + atlas_xaod_dataset() + .Select( + lambda e: { + "FirstJetPt": e.Jets("AntiKt4EMTopoJets") + .Select(lambda j: j.pt() / 1000.0) + .Where(lambda jpt: jpt > 10.0) + .First() + } + ) .value() + ) # Check to see if there mention of push_back anywhere. lines = get_lines_of_code(r) print_lines(lines) @@ -62,9 +74,16 @@ def test_First_Of_Select_is_not_array(): def test_First_Of_Select_After_Where_is_in_right_place(): # Make sure that we have the "First" predicate after if Where's if statement. - r = atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()/1000.0).Where(lambda jpt: jpt > 10.0).First()') \ + r = ( + atlas_xaod_dataset() + .Select( + lambda e: e.Jets("AntiKt4EMTopoJets") + .Select(lambda j: j.pt() / 1000.0) + .Where(lambda jpt: jpt > 10.0) + .First() + ) .value() + ) lines = get_lines_of_code(r) print_lines(lines) ln = find_line_with(">10.0", lines) @@ -73,13 +92,12 @@ def test_First_Of_Select_After_Where_is_in_right_place(): def test_First_with_dict(): - r = (atlas_xaod_dataset() - .Select(lambda e: e.Jets('Anti').First()) - .Select(lambda j: { - 'pt': j.pt(), - 'eta': j.eta() - }) - .value()) + r = ( + atlas_xaod_dataset() + .Select(lambda e: e.Jets("Anti").First()) + .Select(lambda j: {"pt": j.pt(), "eta": j.eta()}) + .value() + ) lines = get_lines_of_code(r) print_lines(lines) diff --git a/tests/atlas/xaod/test_integrated_query.py b/tests/atlas/xaod/test_integrated_query.py index 6b533b53..78a876cc 100644 --- a/tests/atlas/xaod/test_integrated_query.py +++ b/tests/atlas/xaod/test_integrated_query.py @@ -77,8 +77,8 @@ async def test_two_simultaneous_runs(): 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"] - 52.02462890625) < 0.001 # type: ignore () diff --git a/tests/atlas/xaod/test_simple_type_info.py b/tests/atlas/xaod/test_simple_type_info.py index 5cf0e2b8..ca2ad4d4 100644 --- a/tests/atlas/xaod/test_simple_type_info.py +++ b/tests/atlas/xaod/test_simple_type_info.py @@ -8,27 +8,38 @@ def test_cant_call_double(): msg = "" try: - atlas_xaod_dataset("file://root.root") \ - .Select("lambda e: e.Jets('AntiKt4EMTopoJets').Select(lambda j: j.pt().eta()).Sum()") \ - .value() + atlas_xaod_dataset("file://root.root").Select( + lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt().eta()).Sum() + ).value() except xAODTranslationError as e: msg = str(e) - assert 'Unable to call method eta on type double' in msg + assert "Unable to call method eta on type double" in msg def test_can_call_prodVtx(): - ctyp.add_method_type_info("xAOD::TruthParticle", "prodVtx", ctyp.terminal('xAODTruth::TruthVertex', p_depth=1)) - atlas_xaod_dataset("file://root.root") \ - .Select("lambda e: e.TruthParticles('TruthParticles').Select(lambda t: t.prodVtx().x()).Sum()") \ - .value() + ctyp.add_method_type_info( + "xAOD::TruthParticle", + "prodVtx", + ctyp.terminal("xAODTruth::TruthVertex", p_depth=1), + ) + atlas_xaod_dataset("file://root.root").Select( + lambda e: e.TruthParticles("TruthParticles") + .Select(lambda t: t.prodVtx().x()) + .Sum() + ).value() def test_collection_return(): - 'Make sure we can set and deal with a returned collection properly' - ctyp.add_method_type_info("xAOD::TruthParticle", "vertexes", ctyp.collection(ctyp.terminal('double'), p_depth=1)) - (atlas_xaod_dataset("file://root.root") - .SelectMany(lambda e: e.TruthParticles('TruthParticles')) - .SelectMany(lambda t: t.vertexes()) - .value() - ) + "Make sure we can set and deal with a returned collection properly" + ctyp.add_method_type_info( + "xAOD::TruthParticle", + "vertexes", + ctyp.collection(ctyp.terminal("double"), p_depth=1), + ) + ( + atlas_xaod_dataset("file://root.root") + .SelectMany(lambda e: e.TruthParticles("TruthParticles")) + .SelectMany(lambda t: t.vertexes()) + .value() + ) diff --git a/tests/atlas/xaod/test_xaod_aggragate.py b/tests/atlas/xaod/test_xaod_aggragate.py index 895eb957..816d99c6 100644 --- a/tests/atlas/xaod/test_xaod_aggragate.py +++ b/tests/atlas/xaod/test_xaod_aggragate.py @@ -1,4 +1,9 @@ -from tests.utils.locators import find_line_numbers_with, find_line_with, find_next_closing_bracket, find_open_blocks +from tests.utils.locators import ( + find_line_numbers_with, + find_line_with, + find_next_closing_bracket, + find_open_blocks, +) from tests.utils.general import get_lines_of_code, print_lines from tests.atlas.xaod.utils import atlas_xaod_dataset @@ -7,10 +12,14 @@ def test_tree_name(): - r = atlas_xaod_dataset() \ - .Select("lambda e: e.Jets('AntiKt4EMTopoJets').Select(lambda j: j.pt()/1000).Sum()") \ - .AsROOTTTree('junk.root', 'analysis', ['fork']) \ + r = ( + atlas_xaod_dataset() + .Select( + lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt() / 1000).Sum() + ) + .AsROOTTTree("junk.root", "analysis", ["fork"]) .value() + ) lines = get_lines_of_code(r) print_lines(lines) l_sets = find_line_numbers_with('tree("analysis")', lines) @@ -18,9 +27,13 @@ def test_tree_name(): def test_Aggregate_not_initial_const_SUM(): - r = atlas_xaod_dataset() \ - .Select("lambda e: e.Jets('AntiKt4EMTopoJets').Select(lambda j: j.pt()/1000).Sum()") \ + r = ( + atlas_xaod_dataset() + .Select( + lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt() / 1000).Sum() + ) .value() + ) lines = get_lines_of_code(r) print_lines(lines) l_sets = find_line_numbers_with("/1000", lines) @@ -28,19 +41,25 @@ def test_Aggregate_not_initial_const_SUM(): def test_Aggregate_uses_floats_for_float_sum(): - r = atlas_xaod_dataset() \ - .Select("lambda e: e.Jets('AntiKt4EMTopoJets').Select(lambda j: j.pt()/1000).Sum()") \ + r = ( + atlas_xaod_dataset() + .Select( + lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt() / 1000).Sum() + ) .value() + ) lines = get_lines_of_code(r) print_lines(lines) - l_agg_decl = find_line_with('double agg', lines) + l_agg_decl = find_line_with("double agg", lines) assert l_agg_decl > 0 def test_count_after_single_sequence(): - r = atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("AllMyJets").Select(lambda j: j.pt()).Count()') \ + r = ( + atlas_xaod_dataset() + .Select(lambda e: e.Jets("AllMyJets").Select(lambda j: j.pt()).Count()) .value() + ) lines = get_lines_of_code(r) print_lines(lines) # Make sure there is just one for loop in here. @@ -53,9 +72,16 @@ def test_count_after_single_sequence(): def test_count_after_single_sequence_with_filter(): - r = atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("AllMyJets").Select(lambda j: j.pt()).Where(lambda jpt: jpt>10.0).Count()') \ + r = ( + atlas_xaod_dataset() + .Select( + lambda e: e.Jets("AllMyJets") + .Select(lambda j: j.pt()) + .Where(lambda jpt: jpt > 10.0) + .Count() + ) .value() + ) lines = get_lines_of_code(r) print_lines(lines) # Make sure there is just one for loop in here. @@ -68,9 +94,15 @@ def test_count_after_single_sequence_with_filter(): def test_count_after_double_sequence(): - r = atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("AllMyJets").SelectMany(lambda j: e.Tracks("InnerTracks")).Count()') \ + r = ( + atlas_xaod_dataset() + .Select( + lambda e: e.Jets("AllMyJets") + .SelectMany(lambda j: e.Tracks("InnerTracks")) + .Count() + ) .value() + ) lines = get_lines_of_code(r) print_lines(lines) # Make sure there is just one for loop in here. @@ -83,9 +115,15 @@ def test_count_after_double_sequence(): def test_count_after_single_sequence_of_sequence(): - r = atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("AllMyJets").Select(lambda j: e.Tracks("InnerTracks")).Count()') \ + r = ( + atlas_xaod_dataset() + .Select( + lambda e: e.Jets("AllMyJets") + .Select(lambda j: e.Tracks("InnerTracks")) + .Count() + ) .value() + ) lines = get_lines_of_code(r) print_lines(lines) # Make sure there is just one for loop in here. @@ -98,9 +136,17 @@ def test_count_after_single_sequence_of_sequence(): def test_count_after_double_sequence_with_filter(): - r = atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("AllMyJets").SelectMany(lambda j: e.Tracks("InnerTracks").Where(lambda t: t.pt()>10.0)).Count()') \ + r = ( + atlas_xaod_dataset() + .Select( + lambda e: e.Jets("AllMyJets") + .SelectMany( + lambda j: e.Tracks("InnerTracks").Where(lambda t: t.pt() > 10.0) + ) + .Count() + ) .value() + ) lines = get_lines_of_code(r) print_lines(lines) # Make sure there is just one for loop in here. @@ -113,9 +159,16 @@ def test_count_after_double_sequence_with_filter(): def test_count_after_single_sequence_of_sequence_unwound(): - r = atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("AllMyJets").Select(lambda j: e.Tracks("InnerTracks")).SelectMany(lambda ts: ts).Count()') \ + r = ( + atlas_xaod_dataset() + .Select( + lambda e: e.Jets("AllMyJets") + .Select(lambda j: e.Tracks("InnerTracks")) + .SelectMany(lambda ts: ts) + .Count() + ) .value() + ) lines = get_lines_of_code(r) print_lines(lines) # Make sure there is just one for loop in here. @@ -128,13 +181,19 @@ def test_count_after_single_sequence_of_sequence_unwound(): def test_count_after_single_sequence_of_sequence_with_useless_where(): - r = atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("AllMyJets").Select(lambda j: e.Tracks("InnerTracks").Where(lambda pt: pt > 10.0)).Count()') \ + r = ( + atlas_xaod_dataset() + .Select( + lambda e: e.Jets("AllMyJets") + .Select(lambda j: e.Tracks("InnerTracks").Where(lambda pt: pt > 10.0)) + .Count() + ) .value() + ) lines = get_lines_of_code(r) print_lines(lines) # Make sure there is just one for loop in here. - l_increment = find_line_with('+1', lines) + l_increment = find_line_with("+1", lines) block_headers = find_open_blocks(lines[:l_increment]) assert 1 == ["for" in ln for ln in block_headers].count(True) # Make sure the +1 happens after the for, and before another } bracket. @@ -147,48 +206,65 @@ def test_count_after_single_sequence_of_sequence_with_useless_where(): def test_first_can_be_iterable_after_where(): # This was found while trying to generate a tuple for some training, below, simplified. # The problem was that First() always returned something you weren't allowed to iterate over. Which is not what we want here. - atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("AllMyJets").Select(lambda j: e.Tracks("InnerTracks").Where(lambda t: t.pt() > 1000.0)).First().Count()') \ - .value() + atlas_xaod_dataset().Select( + lambda e: e.Jets("AllMyJets") + .Select(lambda j: e.Tracks("InnerTracks").Where(lambda t: t.pt() > 1000.0)) + .First() + .Count() + ).value() def test_first_can_be_iterable(): # Make sure a First() here gets called back correctly and generated. - atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("AllMyJets").Select(lambda j: e.Tracks("InnerTracks")).First().Count()') \ - .value() + atlas_xaod_dataset().Select( + lambda e: e.Jets("AllMyJets") + .Select(lambda j: e.Tracks("InnerTracks")) + .First() + .Count() + ).value() def test_Aggregate_per_jet(): - atlas_xaod_dataset() \ - .Select("lambda e: e.Jets('AntiKt4EMTopoJets').Select(lambda j: j.pt()).Count()") \ - .value() + atlas_xaod_dataset().Select( + lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()).Count() + ).value() def test_Aggregate_per_jet_int(): - r = atlas_xaod_dataset() \ - .Select("lambda e: e.Jets('AntiKt4EMTopoJets').Select(lambda j: j.pt()).Count()") \ + r = ( + atlas_xaod_dataset() + .Select(lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()).Count()) .value() + ) lines = get_lines_of_code(r) print_lines(lines) - l_agg_decl = find_line_with('int agg', lines) + l_agg_decl = find_line_with("int agg", lines) assert l_agg_decl > 0 def test_generate_Max(): - r = atlas_xaod_dataset() \ - .Select("lambda e: e.Jets('AntiKt4EMTopoJets').Select(lambda j: j.pt()).Max()") \ + r = ( + atlas_xaod_dataset() + .Select(lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()).Max()) .value() + ) lines = get_lines_of_code(r) print_lines(lines) def test_First_selects_collection_count(): # Make sure that we have the "First" predicate after if Where's if statement. - r = atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: e.Tracks("InDetTrackParticles")).First().Count()') \ + r = ( + atlas_xaod_dataset() + .Select( + lambda e: e.Jets("AntiKt4EMTopoJets") + .Select(lambda j: e.Tracks("InDetTrackParticles")) + .First() + .Count() + ) .value() + ) lines = get_lines_of_code(r) print_lines(lines) ln = find_line_numbers_with("for", lines) @@ -196,14 +272,25 @@ def test_First_selects_collection_count(): def test_sequence_with_where_first(): - r = atlas_xaod_dataset() \ - .Select('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: e.Tracks("InDetTrackParticles").Where(lambda t: t.pt() > 1000.0)).First().Count()') \ + r = ( + atlas_xaod_dataset() + .Select( + lambda e: e.Jets("AntiKt4EMTopoJets") + .Select( + lambda j: e.Tracks("InDetTrackParticles").Where( + lambda t: t.pt() > 1000.0 + ) + ) + .First() + .Count() + ) .value() + ) lines = get_lines_of_code(r) print_lines(lines) l_first = find_line_numbers_with("if (is_first", lines) assert 1 == len(l_first) - active_blocks = find_open_blocks(lines[:l_first[0]]) + active_blocks = find_open_blocks(lines[: l_first[0]]) assert 1 == ["for" in a for a in active_blocks].count(True) l_agg = find_line_with("+1", lines) active_blocks = find_open_blocks(lines[:l_agg]) diff --git a/tests/atlas/xaod/test_xaod_executor.py b/tests/atlas/xaod/test_xaod_executor.py index 1a4a003d..8fe5d02b 100644 --- a/tests/atlas/xaod/test_xaod_executor.py +++ b/tests/atlas/xaod/test_xaod_executor.py @@ -1,15 +1,15 @@ import re +from math import sin import pytest from func_adl import Range + from func_adl_xAOD.atlas.xaod.executor import atlas_xaod_executor +from func_adl_xAOD.common.math_utils import DeltaR from tests.atlas.xaod.utils import atlas_xaod_dataset, exe_from_qastle # type: ignore from tests.utils.general import get_lines_of_code, print_lines # type: ignore -from tests.utils.locators import ( - find_line_numbers_with, # type: ignore - find_line_with, - find_open_blocks, -) +from tests.utils.locators import find_line_numbers_with # type: ignore +from tests.utils.locators import find_line_with, find_open_blocks # Tests that make sure the xaod executor is working correctly @@ -43,7 +43,7 @@ def test_dict_output_fail_expansion(): def test_per_event_item(): r = ( atlas_xaod_dataset() - .Select('lambda e: e.EventInfo("EventInfo").runNumber()') + .Select(lambda e: e.EventInfo("EventInfo").runNumber()) .value() ) vs = r.QueryVisitor._gc._class_vars @@ -55,7 +55,7 @@ def test_per_jet_item(): # 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.pt())') + .SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt())) .value() ) # Check to see if there mention of push_back anywhere. @@ -128,9 +128,7 @@ def test_builtin_abs_function(): # The following statement should be a straight sequence, not an array. r = ( atlas_xaod_dataset() - .SelectMany( - 'lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: abs(j.pt()))' - ) + .SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: abs(j.pt()))) .value() ) # Check to see if there mention of push_back anywhere. @@ -144,9 +142,7 @@ def test_builtin_sin_function_no_math_import(): # The following statement should be a straight sequence, not an array. r = ( atlas_xaod_dataset() - .SelectMany( - 'lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: sin(j.pt()))' - ) + .SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: sin(j.pt()))) .value() ) # Check to see if there mention of push_back anywhere. @@ -176,7 +172,9 @@ def test_ifexpr(): r = ( atlas_xaod_dataset(qastle_roundtrip=True) .SelectMany( - 'lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: 1.0 if j.pt() > 10.0 else 2.0)' + lambda e: e.Jets("AntiKt4EMTopoJets").Select( + lambda j: 1.0 if j.pt() > 10.0 else 2.0 + ) ) .value() ) @@ -192,8 +190,8 @@ def test_per_jet_item_with_where(): # The following statement should be a straight sequence, not an array. r = ( atlas_xaod_dataset() - .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets")') - .Where("lambda j: j.pt()>40.0") + .SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets")) + .Where(lambda j: j.pt() > 40.0) .Select(lambda j: {"JetPts": j.pt()}) # type: ignore .value() ) @@ -208,9 +206,9 @@ def test_and_clause_in_where(): # The following statement should be a straight sequence, not an array. r = ( atlas_xaod_dataset() - .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets")') - .Where("lambda j: j.pt()>40.0 and j.eta()<2.5") - .Select("lambda j: j.pt()") + .SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets")) + .Where(lambda j: j.pt() > 40.0 and j.eta() < 2.5) + .Select(lambda j: j.pt()) .value() ) # Make sure that the tree Fill is at the same level as the _JetPts2 getting set. @@ -225,9 +223,9 @@ def test_or_clause_in_where(): # The following statement should be a straight sequence, not an array. r = ( atlas_xaod_dataset() - .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets")') - .Where("lambda j: j.pt()>40.0 or j.eta()<2.5") - .Select("lambda j: j.pt()") + .SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets")) + .Where(lambda j: j.pt() > 40.0 or j.eta() < 2.5) + .Select(lambda j: j.pt()) .value() ) # Make sure that the tree Fill is at the same level as the _JetPts2 getting set. @@ -243,8 +241,8 @@ def test_nested_lambda_argument_name_with_monad(): # Need both the monad and the "e" reused to get this error! r = ( atlas_xaod_dataset() - .Select('lambda e: (e.Electrons("Electrons"), e.Muons("Muons"))') - .Select("lambda e: e[0].Select(lambda e: e.E())") + .Select(lambda e: (e.Electrons("Electrons"), e.Muons("Muons"))) + .Select(lambda e: e[0].Select(lambda e: e.E())) .value() ) lines = get_lines_of_code(r) @@ -260,7 +258,7 @@ def test_dict_simple_reference(): .Select( lambda e: {"e_list": e.Electrons("Electrons"), "m_list": e.Muons("Muons")} ) - .Select("lambda e: e.e_list.Select(lambda e: e.E())") + .Select(lambda e: e.e_list.Select(lambda e: e.E())) .value() ) lines = get_lines_of_code(r) @@ -293,8 +291,8 @@ def test_result_awkward(): # 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.pt()") + .SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets")) + .Select(lambda j: j.pt()) .value() ) # Make sure that the tree Fill is at the same level as the _JetPts2 getting set. @@ -308,7 +306,10 @@ def test_per_jet_item_with_event_level(): r = ( atlas_xaod_dataset() .Select( - 'lambda e: (e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()), e.EventInfo("EventInfo").runNumber())' + lambda e: ( + e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()), + e.EventInfo("EventInfo").runNumber(), + ) ) .SelectMany( lambda ji: ji[0].Select( @@ -331,13 +332,13 @@ def test_per_jet_item_with_event_level(): def test_func_sin_call(): atlas_xaod_dataset().Select( - 'lambda e: sin(e.EventInfo("EventInfo").runNumber())' + lambda e: sin(e.EventInfo("EventInfo").runNumber()) ).value() def test_per_jet_item_as_call(): - atlas_xaod_dataset().SelectMany('lambda e: e.Jets("bogus")').Select( - "lambda j: j.pt()" + atlas_xaod_dataset().SelectMany(lambda e: e.Jets("bogus")).Select( + lambda j: j.pt() ).value() @@ -346,7 +347,9 @@ def test_Select_is_an_array_with_where(): r = ( atlas_xaod_dataset() .Select( - 'lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()/1000.0).Where(lambda jpt: jpt > 10.0)' + lambda e: e.Jets("AntiKt4EMTopoJets") + .Select(lambda j: j.pt() / 1000.0) + .Where(lambda jpt: jpt > 10.0) ) .value() ) @@ -363,7 +366,7 @@ def test_Select_is_an_array(): # The following statement should be a straight sequence, not an array. r = ( atlas_xaod_dataset() - .Select('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt())') + .Select(lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt())) .value() ) # Check to see if there mention of push_back anywhere. @@ -380,7 +383,9 @@ def test_Select_1D_array_with_Where(): r = ( atlas_xaod_dataset() .Select( - 'lambda e: e.Jets("AntiKt4EMTopoJets").Where(lambda j1: j1.pt() > 10).Select(lambda j: j.pt())' + lambda e: e.Jets("AntiKt4EMTopoJets") + .Where(lambda j1: j1.pt() > 10) + .Select(lambda j: j.pt()) ) .value() ) @@ -401,7 +406,7 @@ def test_Select_is_not_an_array(): # 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.pt())') + .SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt())) .value() ) # Check to see if there mention of push_back anywhere. @@ -418,7 +423,10 @@ def test_Select_Multiple_arrays(): r = ( atlas_xaod_dataset() .Select( - 'lambda e: (e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()/1000.0),e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.eta()))' + lambda e: ( + e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt() / 1000.0), + e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.eta()), + ) ) .value() ) @@ -435,9 +443,12 @@ def test_Select_Multiple_arrays_2_step(): # The following statement should be a straight sequence, not an array. r = ( atlas_xaod_dataset() - .Select('lambda e: e.Jets("AntiKt4EMTopoJets")') + .Select(lambda e: e.Jets("AntiKt4EMTopoJets")) .Select( - "lambda jets: (jets.Select(lambda j: j.pt()/1000.0),jets.Select(lambda j: j.eta()))" + lambda jets: ( + jets.Select(lambda j: j.pt() / 1000.0), + jets.Select(lambda j: j.eta()), + ) ) .value() ) @@ -462,7 +473,9 @@ def test_Select_of_2D_array(): r = ( atlas_xaod_dataset() .Select( - 'lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: e.Electrons("Electrons").Select(lambda e: e.pt()))' + lambda e: e.Jets("AntiKt4EMTopoJets").Select( + lambda j: e.Electrons("Electrons").Select(lambda e: e.pt()) + ) ) .value() ) @@ -487,7 +500,11 @@ def test_Select_of_2D_with_where(): r = ( atlas_xaod_dataset() .Select( - 'lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: e.Electrons("Electrons").Where(lambda ele: ele.pt() > 10).Select(lambda e: e.pt()))' + lambda e: e.Jets("AntiKt4EMTopoJets").Select( + lambda j: e.Electrons("Electrons") + .Where(lambda ele: ele.pt() > 10) + .Select(lambda e: e.pt()) + ) ) .value() ) @@ -509,7 +526,11 @@ def test_Select_of_3D_array(): r = ( atlas_xaod_dataset() .Select( - 'lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: e.Electrons("Electrons").Select(lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt())))' + lambda e: e.Jets("AntiKt4EMTopoJets").Select( + lambda j: e.Electrons("Electrons").Select( + lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()) + ) + ) ) .value() ) @@ -530,7 +551,9 @@ def test_Select_of_2D_array_with_tuple(): # at least not yet. Make sure error is reasonable. with pytest.raises(Exception) as e: atlas_xaod_dataset().Select( - 'lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: (j.pt()/1000.0, j.eta()))' + lambda e: e.Jets("AntiKt4EMTopoJets").Select( + lambda j: (j.pt() / 1000.0, j.eta()) + ) ).value() assert "data structures" in str(e.value) @@ -541,7 +564,9 @@ def test_SelectMany_of_tuple_is_not_array(): r = ( atlas_xaod_dataset() .SelectMany( - 'lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: (j.pt()/1000.0, j.eta()))' + lambda e: e.Jets("AntiKt4EMTopoJets").Select( + lambda j: (j.pt() / 1000.0, j.eta()) + ) ) .value() ) @@ -573,7 +598,7 @@ def test_generate_binary_operator_pow(): # Make sure the pow operator works correctly - that it doesn't cause a crash in generation. r = ( atlas_xaod_dataset() - .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()**2)') + .SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt() ** 2)) .value() ) lines = get_lines_of_code(r) @@ -587,7 +612,7 @@ def test_generate_binary_operator_unsupported(): # Make sure an unsupported binary operator triggers an exception with pytest.raises(Exception) as e: atlas_xaod_dataset().SelectMany( - 'lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()//2)' + lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt() // 2) ).value() assert "FloorDiv" in str(e) @@ -612,7 +637,7 @@ def test_generate_unary_not(): r = ( atlas_xaod_dataset() .SelectMany( - 'lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: not (j.pt() > 50.0))' + lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: not (j.pt() > 50.0)) ) .value() ) @@ -626,10 +651,21 @@ def test_per_jet_with_matching(): r = ( atlas_xaod_dataset() .Select( - 'lambda e: (e.Jets("AntiKt4EMTopoJets"),e.TruthParticles("TruthParticles").Where(lambda tp1: tp1.pdgId() == 35))' + lambda e: ( + e.Jets("AntiKt4EMTopoJets"), + e.TruthParticles("TruthParticles").Where(lambda tp1: tp1.pdgId() == 35), + ) ) .SelectMany( - "lambda ev: ev[0].Select(lambda j1: (j1, ev[1].Where(lambda tp2: DeltaR(tp2.eta(), tp2.phi(), j1.eta(), j1.phi()) < 0.4)))" + lambda ev: ev[0].Select( + lambda j1: ( + j1, + ev[1].Where( + lambda tp2: DeltaR(tp2.eta(), tp2.phi(), j1.eta(), j1.phi()) + < 0.4 + ), + ) + ) ) .Select( lambda ji: { @@ -653,10 +689,21 @@ def test_per_jet_with_matching_and_zeros(): r = ( atlas_xaod_dataset() .Select( - 'lambda e: (e.Jets("AntiKt4EMTopoJets"),e.TruthParticles("TruthParticles").Where(lambda tp1: tp1.pdgId() == 35))' + lambda e: ( + e.Jets("AntiKt4EMTopoJets"), + e.TruthParticles("TruthParticles").Where(lambda tp1: tp1.pdgId() == 35), + ) ) .SelectMany( - "lambda ev: ev[0].Select(lambda j1: (j1, ev[1].Where(lambda tp2: DeltaR(tp2.eta(), tp2.phi(), j1.eta(), j1.phi()) < 0.4)))" + lambda ev: ev[0].Select( + lambda j1: ( + j1, + ev[1].Where( + lambda tp2: DeltaR(tp2.eta(), tp2.phi(), j1.eta(), j1.phi()) + < 0.4 + ), + ) + ) ) .Select( lambda ji: { @@ -682,19 +729,32 @@ def test_per_jet_with_Count_matching(): # is missed when calculating the y() component (for some reason). This bug may not be in the executor, but, rather, may # be in the function simplifier. # Also, if the "else" doesn't include a "first" thing, then things seem to work just fine too. - # .Where('lambda jall: jall[0].pt() > 40.0') \ r = ( atlas_xaod_dataset() .Select( - 'lambda e: (e.Jets("AntiKt4EMTopoJets"),e.TruthParticles("TruthParticles").Where(lambda tp1: tp1.pdgId() == 35))' + lambda e: ( + e.Jets("AntiKt4EMTopoJets"), + e.TruthParticles("TruthParticles").Where(lambda tp1: tp1.pdgId() == 35), + ) ) .SelectMany( - "lambda ev: ev[0].Select(lambda j1: (j1, ev[1].Where(lambda tp2: DeltaR(tp2.eta(), tp2.phi(), j1.eta(), j1.phi()) < 0.4)))" + lambda ev: ev[0].Select( + lambda j1: ( + j1, + ev[1].Where( + lambda tp2: DeltaR(tp2.eta(), tp2.phi(), j1.eta(), j1.phi()) + < 0.4 + ), + ) + ) ) .Select( - "lambda ji: (ji[0].pt(), 0 if ji[1].Count()==0 else ji[1].First().prodVtx().y())" + lambda ji: ( + ji[0].pt(), + 0 if ji[1].Count() == 0 else ji[1].First().prodVtx().y(), + ) ) - .Where("lambda jall: jall[0] > 40.0") + .Where(lambda jall: jall[0] > 40.0) .value() ) lines = get_lines_of_code(r) @@ -708,15 +768,31 @@ def test_per_jet_with_delta(): r = ( atlas_xaod_dataset() .Select( - 'lambda e: (e.Jets("AntiKt4EMTopoJets"),e.TruthParticles("TruthParticles").Where(lambda tp1: tp1.pdgId() == 35))' + lambda e: ( + e.Jets("AntiKt4EMTopoJets"), + e.TruthParticles("TruthParticles").Where(lambda tp1: tp1.pdgId() == 35), + ) ) .SelectMany( - "lambda ev: ev[0].Select(lambda j1: (j1, ev[1].Where(lambda tp2: DeltaR(tp2.eta(), tp2.phi(), j1.eta(), j1.phi()) < 0.4)))" + lambda ev: ev[0].Select( + lambda j1: ( + j1, + ev[1].Where( + lambda tp2: DeltaR(tp2.eta(), tp2.phi(), j1.eta(), j1.phi()) + < 0.4 + ), + ) + ) ) .Select( - "lambda ji: (ji[0].pt(), 0 if ji[1].Count()==0 else abs(ji[1].First().prodVtx().x()-ji[1].First().decayVtx().x()))" + lambda ji: ( + ji[0].pt(), + 0 + if ji[1].Count() == 0 + else abs(ji[1].First().prodVtx().x() - ji[1].First().decayVtx().x()), + ) ) - .Where("lambda jall: jall[0] > 40.0") + .Where(lambda jall: jall[0] > 40.0) .value() ) lines = get_lines_of_code(r) @@ -731,10 +807,21 @@ def test_per_jet_with_matching_and_zeros_and_sum(): r = ( atlas_xaod_dataset() .Select( - 'lambda e: (e.Jets("AntiKt4EMTopoJets"),e.TruthParticles("TruthParticles").Where(lambda tp1: tp1.pdgId() == 35))' + lambda e: ( + e.Jets("AntiKt4EMTopoJets"), + e.TruthParticles("TruthParticles").Where(lambda tp1: tp1.pdgId() == 35), + ) ) .SelectMany( - "lambda ev: ev[0].Select(lambda j1: (j1, ev[1].Where(lambda tp2: DeltaR(tp2.eta(), tp2.phi(), j1.eta(), j1.phi()) < 0.4)))" + lambda ev: ev[0].Select( + lambda j1: ( + j1, + ev[1].Where( + lambda tp2: DeltaR(tp2.eta(), tp2.phi(), j1.eta(), j1.phi()) + < 0.4 + ), + ) + ) ) .Select( lambda ji: { @@ -759,9 +846,18 @@ def test_electron_and_muon_with_tuple(): # Marc's long query. r = ( atlas_xaod_dataset() - .Select('lambda e: (e.Electrons("Electrons"), e.Muons("Muons"))') + .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()))" + 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()), + ) ) .value() ) @@ -775,9 +871,18 @@ def test_electron_and_muon_with_tuple_qastle(): # Marc's long query. r = ( atlas_xaod_dataset(qastle_roundtrip=True) - .Select('lambda e: (e.Electrons("Electrons"), e.Muons("Muons"))') + .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()))" + 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()), + ) ) .value() ) @@ -791,9 +896,18 @@ def test_electron_and_muon_with_list(): # Marc's long query. r = ( atlas_xaod_dataset() - .Select('lambda e: [e.Electrons("Electrons"), e.Muons("Muons")]') + .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())]" + 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()), + ] ) .value() ) @@ -807,9 +921,18 @@ def test_electron_and_muon_with_list_qastle(): # Marc's long query. r = ( atlas_xaod_dataset(qastle_roundtrip=True) - .Select('lambda e: [e.Electrons("Electrons"), e.Muons("Muons")]') + .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())]" + 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()), + ] ) .value() ) @@ -1027,8 +1150,8 @@ def test_metadata_returned_collection_double_ptr(): def test_executor_forgets_blocks(): "An executor must be able to run twice, and forget between" - from tests.utils.base import dataset, dummy_executor from func_adl_xAOD.atlas.xaod.query_ast_visitor import atlas_xaod_query_ast_visitor + from tests.utils.base import dataset, dummy_executor our_exe = atlas_xaod_executor() diff --git a/tests/cms/aod/test_integrated_query.py b/tests/cms/aod/test_integrated_query.py index 6dde44d1..3b2ddbd0 100644 --- a/tests/cms/aod/test_integrated_query.py +++ b/tests/cms/aod/test_integrated_query.py @@ -11,7 +11,7 @@ pytestmark = run_long_running_tests -if os.name == 'nt': +if os.name == "nt": asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy()) @@ -23,78 +23,87 @@ def turn_on_logging(): def test_select_pt_of_global_muons(): - training_df = as_pandas(f_single - .SelectMany('lambda e: e.TrackMuons("globalMuons")') - .Select('lambda m: m.pt()')) + training_df = as_pandas( + f_single.SelectMany(lambda e: e.TrackMuons("globalMuons")).Select( + lambda m: m.pt() + ) + ) - assert training_df.iloc[0]['col1'] == 10.523032870843785 - assert training_df.iloc[1]['col1'] == 3.914596361447311 - assert training_df.iloc[-1]['col1'] == 29.390803160094887 + assert training_df.iloc[0]["col1"] == 10.523032870843785 + assert training_df.iloc[1]["col1"] == 3.914596361447311 + assert training_df.iloc[-1]["col1"] == 29.390803160094887 def test_select_twice_pt_of_global_muons(): - training_df = as_pandas(f_single - .SelectMany('lambda e: e.TrackMuons("globalMuons")') - .Select('lambda m: m.pt() * 2')) + training_df = as_pandas( + f_single.SelectMany(lambda e: e.TrackMuons("globalMuons")).Select( + lambda m: m.pt() * 2 + ) + ) - assert training_df.iloc[0]['col1'] == 21.04606574168757 - assert training_df.iloc[1]['col1'] == 7.829192722894622 - assert training_df.iloc[-1]['col1'] == 58.78160632018977 + assert training_df.iloc[0]["col1"] == 21.04606574168757 + assert training_df.iloc[1]["col1"] == 7.829192722894622 + assert training_df.iloc[-1]["col1"] == 58.78160632018977 def test_select_eta_of_global_muons(): - training_df = as_pandas(f_single - .SelectMany('lambda e: e.TrackMuons("globalMuons")') - .Select('lambda m: m.eta()')) + training_df = as_pandas( + f_single.SelectMany(lambda e: e.TrackMuons("globalMuons")).Select( + lambda m: m.eta() + ) + ) - assert training_df.iloc[0]['col1'] == -1.8779354371325043 - assert training_df.iloc[1]['col1'] == -2.127157548674547 - assert training_df.iloc[-1]['col1'] == 0.29603168003675756 + assert training_df.iloc[0]["col1"] == -1.8779354371325043 + assert training_df.iloc[1]["col1"] == -2.127157548674547 + assert training_df.iloc[-1]["col1"] == 0.29603168003675756 def test_select_pt_eta_of_global_muons(): - training_df = as_pandas(f_single - .SelectMany('lambda e: e.TrackMuons("globalMuons")') - .Select('lambda m: m.pt() + m.eta()')) + training_df = as_pandas( + f_single.SelectMany(lambda e: e.TrackMuons("globalMuons")).Select( + lambda m: m.pt() + m.eta() + ) + ) - assert training_df.iloc[0]['col1'] == 8.645097433711282 - assert training_df.iloc[1]['col1'] == 1.787438812772764 - assert training_df.iloc[-1]['col1'] == 29.686834840131645 + assert training_df.iloc[0]["col1"] == 8.645097433711282 + assert training_df.iloc[1]["col1"] == 1.787438812772764 + assert training_df.iloc[-1]["col1"] == 29.686834840131645 def test_select_hitpattern_of_global_muons(): sys.setrecursionlimit(10000) - training_df = as_pandas(f_single - .SelectMany(lambda e: e.TrackMuons("globalMuons")) - .Select(lambda m: m.hitPattern()) - .Select(lambda hp: Range(0, hp.numberOfHits()) - .Select(lambda i: hp.getHitPattern(i)))) + training_df = as_pandas( + f_single.SelectMany(lambda e: e.TrackMuons("globalMuons")) + .Select(lambda m: m.hitPattern()) + .Select( + lambda hp: Range(0, hp.numberOfHits()).Select(lambda i: hp.getHitPattern(i)) + ) + ) - assert training_df.iloc[0]['col1'] == 1160.0 - assert training_df.iloc[1]['col1'] == 1168.0 - assert training_df.iloc[-1]['col1'] == 248.0 + assert training_df.iloc[0]["col1"] == 1160.0 + assert training_df.iloc[1]["col1"] == 1168.0 + assert training_df.iloc[-1]["col1"] == 248.0 def test_isnonull_call(): - ''' Make sure the non null call works properly. This is tricky because of the + """Make sure the non null call works properly. This is tricky because of the way objects are used here - both as a pointer and an object, so the code has to work just right. - ''' + """ training_df = as_pandas( - f_single - .Select(lambda e: e.Muons("muons")) + f_single.Select(lambda e: e.Muons("muons")) .Select(lambda muons: muons.Where(lambda m: isNonnull(m.globalTrack()))) .Select(lambda muons: muons.Count()) ) - assert training_df.iloc[0]['col1'] == 3 + assert training_df.iloc[0]["col1"] == 3 def test_sumChargedHadronPt(): training_df = as_pandas( - f_single - .SelectMany(lambda e: e.Muons("muons")) - .Select(lambda m: (m.pfIsolationR04()).sumChargedHadronPt) + f_single.SelectMany(lambda e: e.Muons("muons")).Select( + lambda m: (m.pfIsolationR04()).sumChargedHadronPt + ) ) - assert training_df.iloc[0]['col1'] == 1.6184179782867432 - assert training_df.iloc[1]['col1'] == 0.8745666742324829 - assert training_df.iloc[-1]['col1'] == 0.0 + assert training_df.iloc[0]["col1"] == 1.6184179782867432 + assert training_df.iloc[1]["col1"] == 0.8745666742324829 + assert training_df.iloc[-1]["col1"] == 0.0 diff --git a/tests/cms/aod/test_local_dataset.py b/tests/cms/aod/test_local_dataset.py index 9a9fbbd1..3dc34305 100644 --- a/tests/cms/aod/test_local_dataset.py +++ b/tests/cms/aod/test_local_dataset.py @@ -6,44 +6,50 @@ @pytest.mark.cms_aod_runner def test_integrated_run(): - '''Test a simple run with docker''' + """Test a simple run with docker""" # TODO: Using the type stuff, make sure replacing Select below with SelectMany makes a good error message from func_adl_xAOD.cms.aod.local_dataset import CMSRun1AODDataset - r = (CMSRun1AODDataset(f_location) - .SelectMany('lambda e: e.TrackMuons("globalMuons")') - .Select('lambda m: m.pt()') - .AsROOTTTree('junk.root', 'my_tree', ['muon_pt']) - .value()) + + r = ( + CMSRun1AODDataset(f_location) + .SelectMany(lambda e: e.TrackMuons("globalMuons")) + .Select(lambda m: m.pt()) + .AsROOTTTree("junk.root", "my_tree", ["muon_pt"]) + .value() + ) assert len(r) == 1 @pytest.fixture() def docker_mock(mocker): - 'Mock the docker object' + "Mock the docker object" m = mocker.MagicMock(spec=python_on_whales.docker) def parse_arg(*args, **kwargs): - v = kwargs['volumes'] - data_s = [d for d in v if d[1] == '/results'] + v = kwargs["volumes"] + data_s = [d for d in v if d[1] == "/results"] assert len(data_s) == 1 data = data_s[0][0] - (data / 'ANALYSIS.root').touch() + (data / "ANALYSIS.root").touch() - return (('stdout', b'text'), ('stderr', b'text')) + return (("stdout", b"text"), ("stderr", b"text")) m.run.side_effect = parse_arg - mocker.patch('func_adl_xAOD.common.local_dataset.docker', m) + mocker.patch("func_adl_xAOD.common.local_dataset.docker", m) return m def test_run(docker_mock): - '''Test a simple run using docker mock''' + """Test a simple run using docker mock""" from func_adl_xAOD.cms.aod.local_dataset import CMSRun1AODDataset - r = (CMSRun1AODDataset(f_location) - .SelectMany('lambda e: e.TrackMuons("globalMuons")') - .Select('lambda m: m.pt()') - .AsROOTTTree('junk.root', 'my_tree', ['muon_pt']) - .value()) + + r = ( + CMSRun1AODDataset(f_location) + .SelectMany(lambda e: e.TrackMuons("globalMuons")) + .Select(lambda m: m.pt()) + .AsROOTTTree("junk.root", "my_tree", ["muon_pt"]) + .value() + ) assert len(r) == 1 diff --git a/tests/cms/miniaod/test_integrated_query.py b/tests/cms/miniaod/test_integrated_query.py index fb9ef114..e7196c55 100644 --- a/tests/cms/miniaod/test_integrated_query.py +++ b/tests/cms/miniaod/test_integrated_query.py @@ -11,7 +11,7 @@ pytestmark = run_long_running_tests -if os.name == 'nt': +if os.name == "nt": asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy()) @@ -23,89 +23,100 @@ def turn_on_logging(): def test_select_pt_of_muons(): - training_df = as_pandas(f_single - .SelectMany('lambda e: e.Muons("slimmedMuons")') - .Select('lambda m: m.pt()')) + training_df = as_pandas( + f_single.SelectMany(lambda e: e.Muons("slimmedMuons")).Select(lambda m: m.pt()) + ) - assert training_df.iloc[0]['col1'] == 12.901309967041016 - assert training_df.iloc[1]['col1'] == 7.056070327758789 - assert training_df.iloc[-1]['col1'] == 49.50603485107422 + assert training_df.iloc[0]["col1"] == 12.901309967041016 + assert training_df.iloc[1]["col1"] == 7.056070327758789 + assert training_df.iloc[-1]["col1"] == 49.50603485107422 def test_select_pt_of_electrons(): - training_df = as_pandas(f_single - .SelectMany('lambda e: e.Electrons("slimmedElectrons")') - .Select('lambda m: m.pt()')) + training_df = as_pandas( + f_single.SelectMany(lambda e: e.Electrons("slimmedElectrons")).Select( + lambda m: m.pt() + ) + ) - assert training_df.iloc[0]['col1'] == 4.862127780914307 - assert training_df.iloc[1]['col1'] == 21.7811222076416 - assert training_df.iloc[-1]['col1'] == 7.3344407081604 + assert training_df.iloc[0]["col1"] == 4.862127780914307 + assert training_df.iloc[1]["col1"] == 21.7811222076416 + assert training_df.iloc[-1]["col1"] == 7.3344407081604 def test_select_twice_pt_of_global_muons(): - training_df = as_pandas(f_single - .SelectMany('lambda e: e.Muons("slimmedMuons")') - .Select('lambda m: m.pt() * 2')) + training_df = as_pandas( + f_single.SelectMany(lambda e: e.Muons("slimmedMuons")).Select( + lambda m: m.pt() * 2 + ) + ) - assert training_df.iloc[0]['col1'] == 25.80261993408203 - assert training_df.iloc[1]['col1'] == 14.112140655517578 - assert training_df.iloc[-1]['col1'] == 99.01206970214844 + assert training_df.iloc[0]["col1"] == 25.80261993408203 + assert training_df.iloc[1]["col1"] == 14.112140655517578 + assert training_df.iloc[-1]["col1"] == 99.01206970214844 def test_select_eta_of_global_muons(): - training_df = as_pandas(f_single - .SelectMany('lambda e: e.Muons("slimmedMuons")') - .Select('lambda m: m.eta()')) + training_df = as_pandas( + f_single.SelectMany(lambda e: e.Muons("slimmedMuons")).Select(lambda m: m.eta()) + ) - assert training_df.iloc[0]['col1'] == -1.2982407808303833 - assert training_df.iloc[1]['col1'] == -1.1389755010604858 - assert training_df.iloc[-1]['col1'] == 1.2762727737426758 + assert training_df.iloc[0]["col1"] == -1.2982407808303833 + assert training_df.iloc[1]["col1"] == -1.1389755010604858 + assert training_df.iloc[-1]["col1"] == 1.2762727737426758 def test_select_pt_eta_of_global_muons(): - training_df = as_pandas(f_single - .SelectMany('lambda e: e.Muons("slimmedMuons")') - .Select('lambda m: m.pt() + m.eta()')) + training_df = as_pandas( + f_single.SelectMany(lambda e: e.Muons("slimmedMuons")).Select( + lambda m: m.pt() + m.eta() + ) + ) - assert training_df.iloc[0]['col1'] == 11.603069186210632 - assert training_df.iloc[1]['col1'] == 5.917094826698303 - assert training_df.iloc[-1]['col1'] == 50.782307624816895 + assert training_df.iloc[0]["col1"] == 11.603069186210632 + assert training_df.iloc[1]["col1"] == 5.917094826698303 + assert training_df.iloc[-1]["col1"] == 50.782307624816895 def test_select_hitpattern_of_global_muons(): sys.setrecursionlimit(10000) - training_df = as_pandas(f_single - .SelectMany(lambda e: e.Muons("slimmedMuons") - .Where(lambda m: isNonnull(m.globalTrack())) - .Select(lambda m: m.globalTrack()) - .Select(lambda m: m.hitPattern()) - .Select(lambda hp: Range(0, hp.numberOfHits(hp.TRACK_HITS)) - .Select(lambda i: hp.getHitPattern(hp.TRACK_HITS, i))))) - assert training_df.iloc[0]['col1'] == 1416.0 - assert training_df.iloc[1]['col1'] == 1420.0 - assert training_df.iloc[-1]['col1'] == 488.0 + training_df = as_pandas( + f_single.SelectMany( + lambda e: e.Muons("slimmedMuons") + .Where(lambda m: isNonnull(m.globalTrack())) + .Select(lambda m: m.globalTrack()) + .Select(lambda m: m.hitPattern()) + .Select( + lambda hp: Range(0, hp.numberOfHits(hp.TRACK_HITS)).Select( + lambda i: hp.getHitPattern(hp.TRACK_HITS, i) + ) + ) + ) + ) + assert training_df.iloc[0]["col1"] == 1416.0 + assert training_df.iloc[1]["col1"] == 1420.0 + assert training_df.iloc[-1]["col1"] == 488.0 def test_isnonull_call(): - ''' Make sure the non null call works properly. This is tricky because of the + """Make sure the non null call works properly. This is tricky because of the way objects are used here - both as a pointer and an object, so the code has to work just right. - ''' + """ training_df = as_pandas( - f_single - .Select(lambda e: e.Muons("slimmedMuons")) + f_single.Select(lambda e: e.Muons("slimmedMuons")) .Select(lambda muons: muons.Where(lambda m: isNonnull(m.globalTrack()))) .Select(lambda muons: muons.Count()) ) - assert training_df.iloc[0]['col1'] == 1 + assert training_df.iloc[0]["col1"] == 1 def test_sumChargedHadronPt(): training_df = as_pandas( - f_single - .SelectMany(lambda e: e.Muons("slimmedMuons")) - .Select(lambda m: (m.pfIsolationR04()).sumChargedHadronPt) + f_single.SelectMany(lambda e: e.Muons("slimmedMuons")).Select( + lambda m: (m.pfIsolationR04()).sumChargedHadronPt + ) ) - assert training_df.iloc[0]['col1'] == 0.0 - assert training_df.iloc[2]['col1'] == 26.135541915893555 - assert training_df.iloc[-1]['col1'] == 38.0556526184082 + assert training_df.iloc[0]["col1"] == 0.0 + assert training_df.iloc[2]["col1"] == 26.135541915893555 + assert training_df.iloc[-1]["col1"] == 38.0556526184082 diff --git a/tests/cms/miniaod/test_local_dataset.py b/tests/cms/miniaod/test_local_dataset.py index 3a3a82b8..62b8d1cd 100644 --- a/tests/cms/miniaod/test_local_dataset.py +++ b/tests/cms/miniaod/test_local_dataset.py @@ -6,44 +6,50 @@ @pytest.mark.cms_miniaod_runner def test_integrated_run(): - '''Test a simple run with docker''' + """Test a simple run with docker""" # TODO: Using the type stuff, make sure replacing Select below with SelectMany makes a good error message from func_adl_xAOD.cms.miniaod.local_dataset import CMSRun2miniAODDataset - r = (CMSRun2miniAODDataset(f_location) - .SelectMany('lambda e: e.Muons("slimmedMuons")') - .Select('lambda m: m.pt()') - .AsROOTTTree('junk.root', 'my_tree', ['muon_pt']) - .value()) + + r = ( + CMSRun2miniAODDataset(f_location) + .SelectMany(lambda e: e.Muons("slimmedMuons")) + .Select(lambda m: m.pt()) + .AsROOTTTree("junk.root", "my_tree", ["muon_pt"]) + .value() + ) assert len(r) == 1 @pytest.fixture() def docker_mock(mocker): - 'Mock the docker object' + "Mock the docker object" m = mocker.MagicMock(spec=python_on_whales.docker) def parse_arg(*args, **kwargs): - v = kwargs['volumes'] - data_s = [d for d in v if d[1] == '/results'] + v = kwargs["volumes"] + data_s = [d for d in v if d[1] == "/results"] assert len(data_s) == 1 data = data_s[0][0] - (data / 'ANALYSIS.root').touch() + (data / "ANALYSIS.root").touch() - return (('stdout', b'text'), ('stderr', b'text')) + return (("stdout", b"text"), ("stderr", b"text")) m.run.side_effect = parse_arg - mocker.patch('func_adl_xAOD.common.local_dataset.docker', m) + mocker.patch("func_adl_xAOD.common.local_dataset.docker", m) return m def test_run(docker_mock): - '''Test a simple run using docker mock''' + """Test a simple run using docker mock""" from func_adl_xAOD.cms.miniaod.local_dataset import CMSRun2miniAODDataset - r = (CMSRun2miniAODDataset(f_location) - .SelectMany('lambda e: e.Muons("slimmedMuons")') - .Select('lambda m: m.pt()') - .AsROOTTTree('junk.root', 'my_tree', ['muon_pt']) - .value()) + + r = ( + CMSRun2miniAODDataset(f_location) + .SelectMany(lambda e: e.Muons("slimmedMuons")) + .Select(lambda m: m.pt()) + .AsROOTTTree("junk.root", "my_tree", ["muon_pt"]) + .value() + ) assert len(r) == 1