def test_2nd_order_lookup(): 'Seen in the wild to generate an out-of-scope error' r = (cms_aod_dataset().Select( lambda e: { "m": e.Muons("muons"), "p": e.Vertex("offlinePrimaryVertices")[0].position() }).Select( lambda i: i.m.Where(lambda m: m.isPFMuon( ) and m.isPFIsolationValid() and isNonnull(m.globalTrack( )) and abs((m.globalTrack()).dxy(i.p)) < 0.5 and abs( (m.globalTrack()).dz(i.p)) < 1.).Select(lambda m: m.p()), ).value()) lines = get_lines_of_code(r) print_lines(lines) # Make sure the vertex line isn't used after it goes out of scope vertex_decl_line = find_line_with('edm::Handle<reco::VertexCollection>', lines) vertex_variable_name = lines[vertex_decl_line].split(' ')[-1].strip(';') closing_scope = find_next_closing_bracket(lines[vertex_decl_line:]) vertex_used_too_late = find_line_with(vertex_variable_name, lines[vertex_decl_line + closing_scope:], throw_if_not_found=False) if vertex_used_too_late != -1: print('Here is where it is used and down') print_lines(lines[closing_scope + vertex_decl_line + vertex_used_too_late:]) assert vertex_used_too_late == -1
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()') \ .value() lines = get_lines_of_code(r) print_lines(lines) ln = find_line_with(">10.0", lines) # Look for the "false" that First uses to remember it has gone by one. assert find_line_with("false", lines[ln:], throw_if_not_found=False) > 0
def test_Select_member_variable(): r = cms_aod_dataset() \ .SelectMany(lambda e: e.Muons("muons")) \ .Select(lambda m: m.pfIsolationR04().sumChargedHadronPt) \ .value() lines = get_lines_of_code(r) _ = find_line_with(".sumChargedHadronPt", lines) assert find_line_with(".sumChargedHadronPt()", lines, throw_if_not_found=False) == -1
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)') \ .value() lines = get_lines_of_code(r) print_lines(lines) l1 = find_line_with("pow(i_obj", lines) l2 = find_line_with("->pt(), 2)", lines) assert l1 == l2
def test_dict_simple_reference_prop_lookup(): 'Dictionary references should be resolved automatically' r = atlas_xaod_dataset() \ .Select(lambda e: {'e_list': e.Electrons("Electrons"), 'm_list': e.Muons("Muons")}) \ .Select(lambda e: e['e_list'].Select(lambda e: e.E())) \ .value() lines = get_lines_of_code(r) print_lines(lines) l_push = find_line_with('push_back', lines) assert "->E()" in lines[l_push] r = find_line_with('muon', lines, throw_if_not_found=False) assert r == -1
def test_Range_good_call(): # The following statement should be a straight sequence, not an array. r = (atlas_xaod_dataset().SelectMany( lambda e: e.Jets("AntiKt4EMTopoJets")).Select(lambda j: Range( 0, 10).Select(lambda index: j.pt() * index)).value()) # Check to see if there mention of push_back anywhere. lines = get_lines_of_code(r) print_lines(lines) for_loops = find_line_numbers_with('for (', lines) assert len(for_loops) == 2 find_line_with('(0)', lines) find_line_with('(10)', lines)
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()') \ .value() lines = get_lines_of_code(r) print_lines(lines) # Make sure there is just one for loop in here. assert 2 == ["for" in ln for ln in lines].count(True) # Make sure the +1 happens after the for, and before another } bracket. num_for = find_line_with("if", lines) num_inc = find_line_with("+1", lines[num_for:]) num_close = find_next_closing_bracket(lines[num_for:]) assert num_close > num_inc
def test_complex_dict(): 'Seen to fail in the wild, so a test case to track' r = cms_aod_dataset() \ .Select(lambda e: {"muons": e.Muons("muons"), "primvtx": e.Vertex("offlinePrimaryVertices")}) \ .Select(lambda i: i.muons .Where(lambda m: isNonnull(m.globalTrack())) .Select(lambda m: m.globalTrack().dx(i.primvtx[0].position())) ) \ .value() lines = get_lines_of_code(r) print_lines(lines) find_line_with("globalTrack()->dx", lines) find_line_with("at(0).position()", lines)
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())') \ .SelectMany(lambda ji: ji[0].Select(lambda pt: { # type: ignore 'JetPt': pt, 'runNumber': ji[1]} # type: ignore )) \ .value() lines = get_lines_of_code(r) print_lines(lines) l_jetpt = find_line_with("_JetPt", lines) l_runnum = find_line_with("_runNumber", lines) l_fill = find_line_with("->Fill()", lines) assert l_jetpt + 1 == l_runnum assert l_runnum + 1 == l_fill
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()') \ .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) 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. num_for = find_line_with("for", lines) num_inc = find_line_with("+1", lines[num_for:]) num_close = find_next_closing_bracket(lines[num_for:]) assert num_close > num_inc
def test_Select_of_2D_with_where(): # This should generate a 2D array. 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()))') \ .value() lines = get_lines_of_code(r) print_lines(lines) l_vector_decl = find_line_with("vector<double>", lines) l_vector_active = len(find_open_blocks(lines[:l_vector_decl])) l_first_push = find_line_with("push_back", lines) l_first_push_active = len(find_open_blocks(lines[:l_first_push])) assert ( l_vector_active + 2 ) == l_first_push_active # +2 because it is inside the for loop and the if block
def test_Select_of_3D_array(): # This should generate a 2D 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())))') \ .value() lines = get_lines_of_code(r) print_lines(lines) l_vector_decl = find_line_with("vector<double> ", lines) l_vector_active = len(find_open_blocks(lines[:l_vector_decl])) l_vector_double_decl = find_line_with("vector<std::vector<double>>", lines) l_vector_double_active = len(find_open_blocks( lines[:l_vector_double_decl])) assert l_vector_active == (l_vector_double_active + 1)
def test_generate_unary_not(): r = atlas_xaod_dataset() \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: not (j.pt() > 50.0))') \ .value() lines = get_lines_of_code(r) print_lines(lines) _ = find_line_with("!(", lines)
def test_Select_1D_array_with_Where(): # The following statement should be a straight sequence, not an array. r = atlas_xaod_dataset() \ .Select('lambda e: e.Jets("AntiKt4EMTopoJets").Where(lambda j1: j1.pt() > 10).Select(lambda j: j.pt())') \ .value() # Check to see if there mention of push_back anywhere. lines = get_lines_of_code(r) print_lines(lines) assert 1 == ["push_back" in ln for ln in lines].count(True) l_push_back = find_line_with("Fill()", lines) active_blocks = find_open_blocks(lines[:l_push_back]) assert 0 == ["for" in a for a in active_blocks].count(True) push_back = find_line_with("push_back", lines) active_blocks = find_open_blocks(lines[:push_back]) assert 1 == ['if' in a for a in active_blocks].count(True)
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()") \ .value() lines = get_lines_of_code(r) print_lines(lines) l_agg_decl = find_line_with('double agg', lines) assert l_agg_decl > 0
def test_per_jet_with_matching_and_zeros(): # Trying to repro a bug we saw in the wild r = atlas_xaod_dataset() \ .Select('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)))') \ .Select(lambda ji: { 'JetPts': ji[0].pt(), # type: ignore 'NumLLPs': 0 if ji[1].Count() == 0 else (ji[1].First().pt() - ji[1].First().pt()) # type: ignore }) \ .value() lines = get_lines_of_code(r) print_lines(lines) l_jetpt = find_line_with("_JetPts", lines) l_nllp = find_line_with("_NumLLPs", lines) l_fill = find_line_with("->Fill()", lines) assert l_jetpt + 1 == l_nllp assert l_nllp + 1 == l_fill
def test_generate_unary_operations(): ops = ['+', '-'] for o in ops: r = atlas_xaod_dataset() \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()+({0}1))'.format(o)) \ .value() lines = get_lines_of_code(r) print_lines(lines) _ = find_line_with(f"pt()+({o}(1))", lines)
def test_Aggregate_per_jet_int(): 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) assert l_agg_decl > 0
def test_electron_and_muon_with_list_qastle(): # See if we can re-create a bug we are seeing with # Marc's long query. r = atlas_xaod_dataset(qastle_roundtrip=True) \ .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())]') \ .value() lines = get_lines_of_code(r) print_lines(lines) assert find_line_with("->Fill()", lines) != 0
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"))') \ .value() # Check to see if there mention of push_back anywhere. lines = get_lines_of_code(r) print_lines(lines) l_attribute = find_line_with("getAttribute<", lines) assert 'getAttribute<float>("emf")' in lines[l_attribute]
def test_generate_binary_operators(): # Make sure the binary operators work correctly - that they don't cause a crash in generation. ops = ['+', '-', '*', '/', '%'] for o in ops: r = atlas_xaod_dataset() \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt(){0}1)'.format(o)) \ .value() lines = get_lines_of_code(r) print_lines(lines) _ = find_line_with(f"pt(){o}1", lines)
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()))') \ .value() # Check to see if there mention of push_back anywhere. lines = get_lines_of_code(r) print_lines(lines) l_abs = find_line_with("std::sin", lines) assert "->pt()" in lines[l_abs]
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())') \ .value() lines = get_lines_of_code(r) print_lines(lines) l_push = find_line_with('push_back', lines) assert "->E()" in lines[l_push]
def test_SelectMany_of_tuple_is_not_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()/1000.0, j.eta()))') \ .value() lines = get_lines_of_code(r) print_lines(lines) assert 0 == ["push_back" in ln for ln in lines].count(True) l_push_back = find_line_with("Fill()", lines) active_blocks = find_open_blocks(lines[:l_push_back]) assert 1 == ["for" in a for a in active_blocks].count(True)
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()") \ .value() # Make sure that the tree Fill is at the same level as the _JetPts2 getting set. lines = get_lines_of_code(r) print_lines(lines) l_jetpt = find_line_with("_col1", lines) assert "Fill()" in lines[l_jetpt + 1]
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() }) \ .value() # Check to see if there mention of push_back anywhere. lines = get_lines_of_code(r) print_lines(lines) assert all("push_back" not in ln for ln in lines) l_fill = find_line_with("Fill()", lines) active_blocks = find_open_blocks(lines[:l_fill]) assert 3 == [(("for" in a) or ("if" in a)) for a in active_blocks].count(True) l_set = find_line_with("_FirstJetPt", lines) active_blocks = find_open_blocks(lines[:l_set]) assert 3 == [(("for" in a) or ("if" in a)) for a in active_blocks].count(True) l_true = find_line_with("(true)", lines) active_blocks = find_open_blocks(lines[:l_true]) assert 0 == [(("for" in a) or ("if" in a)) for a in active_blocks].count(True)
def test_per_jet_dict_items(): # The following statement should be a straight sequence, not an array. r = (atlas_xaod_dataset().SelectMany( lambda e: e.Jets("AntiKt4EMTopoJets")).Select(lambda j: { 'pt': j.pt(), 'eta': j.eta(), }).value()) # Check to see if there mention of push_back anywhere. lines = get_lines_of_code(r) print_lines(lines) l_pt = lines[find_line_with("->pt()", lines)] l_eta = lines[find_line_with("->eta()", lines)] obj_finder = re.compile(r".*(i_obj[0-9]+)->.*") l_pt_r = obj_finder.match(l_pt) l_eta_r = obj_finder.match(l_eta) assert l_pt_r is not None assert l_eta_r is not None assert l_pt_r[1] == l_eta_r[1]
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()) lines = get_lines_of_code(r) print_lines(lines) l_pt = lines[find_line_with("->pt()", lines)] l_eta = lines[find_line_with("->eta()", lines)] obj_finder = re.compile(r".*(i_obj[0-9]+)->.*") l_pt_r = obj_finder.match(l_pt) l_eta_r = obj_finder.match(l_eta) assert l_pt_r is not None assert l_eta_r is not None assert l_pt_r[1] == l_eta_r[1]
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()') \ .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]]) 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]) assert 1 == [">1000" in a for a in active_blocks].count(True)
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") \ .Select(lambda j: { 'JetPts': j.pt() # type: ignore }) \ .value() # Make sure that the tree Fill is at the same level as the _JetPts2 getting set. lines = get_lines_of_code(r) print_lines(lines) l_jetpt = find_line_with("_JetPts", lines) assert "Fill()" in lines[l_jetpt + 1]