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_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_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_Aggregate_not_initial_const_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)
assert 1 == len(l_sets)
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_tree_name():
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)
assert 1 == len(l_sets)
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()') \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
ln = find_line_numbers_with("for", lines)
assert 2 == len(ln)
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_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_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_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_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_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_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_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)') \
.value()
# Make sure that a test around 10.0 occurs.
lines = get_lines_of_code(r)
print_lines(lines)
lines = [ln for ln in lines if '10.0' in ln]
assert len(lines) == 1
assert 'if ' in lines[0]
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_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_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_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())') \
.value()
# Check to see if there mention of push_back anywhere.
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_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_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_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_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()") \
.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_if = [ln for ln in lines if "if (" in ln]
assert len(l_if) == 2
assert l_if[0] == l_if[1]
def test_per_jet_with_delta():
# 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: (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') \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
l_numbers = find_line_numbers_with("if (i_obj", lines)
for line in [lines[ln] for ln in l_numbers]:
assert "x()" not in line
def test_metadata_returned_type():
# The following statement should be a straight sequence, not an array.
r = (atlas_xaod_dataset().MetaData({
'metadata_type': 'add_method_type_info',
'type_string': 'xAOD::Jet',
'method_name': 'pt',
'return_type': 'double',
}).SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets")).Select(
lambda j: j.pt() * 2).value())
# Check to see if there mention of push_back anywhere.
lines = get_lines_of_code(r)
print_lines(lines)
value_ref = find_line_numbers_with('->pt()*2', lines)
assert len(value_ref) == 1
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_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]
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_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
