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_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_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_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_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 jets: (jets.Select(lambda j: j.pt()/1000.0),jets.Select(lambda j: j.eta()))') \
.value()
# Check to see if there mention of push_back anywhere.
lines = get_lines_of_code(r)
print_lines(lines)
l_push_back = find_line_numbers_with("push_back", lines)
assert all([
len([ln for ln in find_open_blocks(lines[:pb]) if "for" in ln]) == 1
for pb in l_push_back
])
assert 2 == ["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)
def test_Select_of_2D_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.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_numbers_with("push_back", lines)
assert len(l_first_push) == 2
l_first_push_active = len(find_open_blocks(lines[:l_first_push[0]]))
assert (l_vector_active + 1) == l_first_push_active
# Now, make sure the second push_back is at the right level.
l_second_push_active = len(find_open_blocks(lines[:l_first_push[1]]))
assert (l_second_push_active + 1) == l_first_push_active
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_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_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