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