def test_sequence_type():
tc = gc_scope_top_level()
s_value = crep.cpp_value('0.0', tc, ctyp.terminal('int', False))
i_value = crep.cpp_value('1.0', tc, ctyp.terminal('object', False))
seq = crep.cpp_sequence(s_value, i_value, tc)
seq_array = crep.cpp_sequence(seq, i_value, tc)
assert seq_array.sequence_value().cpp_type().type == 'std::vector<int>'
def call_Where(self, node: ast.AST, args: List[ast.AST]):
'Apply a filtering to the current loop.'
assert len(args) == 2
source = args[0]
filter = args[1]
assert isinstance(filter, ast.Lambda)
# Make sure we are in a loop
seq = self.as_sequence(source)
# Simulate the filtering call - we want the resulting value to test.
filter = lambda_unwrap(filter)
c = ast.Call(func=filter, args=[seq.sequence_value().as_ast()])
rep = self.get_rep(c)
# Create an if statement
self._gc.add_statement(statement.iftest(rep))
# Ok - new sequence. This the same as the old sequence, only the sequence value is updated.
# Protect against sequence of sequences (LOVE type checkers, which caught this as a possibility)
w_val = seq.sequence_value()
if isinstance(w_val, crep.cpp_sequence):
raise Exception("Error: A Where clause must evaluate to a value, not a sequence")
new_sequence_var = w_val.copy_with_new_scope(self._gc.current_scope())
node.rep = crep.cpp_sequence(new_sequence_var, seq.iterator_value(), self._gc.current_scope()) # type: ignore
self._result = node.rep # type: ignore
def write_cpp_files(self, ast: ast.AST,
output_path: Path) -> xAODExecutionInfo:
r"""
Given the AST generate the C++ files that need to run. Return them along with
the input files.
"""
# Find the base file dataset and mark it.
from func_adl import find_EventDataset
file = find_EventDataset(ast)
iterator = crep.cpp_variable("bogus-do-not-use",
top_level_scope(),
cpp_type=None)
file.rep = crep.cpp_sequence(iterator, iterator,
top_level_scope()) # type: ignore
# Visit the AST to generate the code structure and find out what the
# result is going to be.
qv = query_ast_visitor()
result_rep = qv.get_rep(ast) if _is_format_request(ast) \
else qv.get_as_ROOT(ast)
# Emit the C++ code into our dictionaries to be used in template generation below.
query_code = _cpp_source_emitter()
qv.emit_query(query_code)
book_code = _cpp_source_emitter()
qv.emit_book(book_code)
class_dec_code = qv.class_declaration_code()
includes = qv.include_files()
# The replacement dict to pass to the template generator can now be filled
info = {}
info['query_code'] = query_code.lines_of_query_code()
info['book_code'] = book_code.lines_of_query_code()
info['class_dec'] = class_dec_code
info['include_files'] = includes
# We use jinja2 templates. Write out everything.
template_dir = _find_dir("func_adl_xAOD/R21Code")
j2_env = jinja2.Environment(
loader=jinja2.FileSystemLoader(template_dir))
self._copy_template_file(j2_env, info, 'ATestRun_eljob.py',
output_path)
self._copy_template_file(j2_env, info, 'package_CMakeLists.txt',
output_path)
self._copy_template_file(j2_env, info, 'query.cxx', output_path)
self._copy_template_file(j2_env, info, 'query.h', output_path)
self._copy_template_file(j2_env, info, 'runner.sh', output_path)
(output_path / 'runner.sh').chmod(0o755)
# Build the return object.
return xAODExecutionInfo(result_rep, output_path, 'runner.sh', [
'ATestRun_eljob.py', 'package_CMakeLists.txt', 'query.cxx',
'query.h', 'runner.sh'
])
def test_as_root_as_single_column():
q = query_ast_visitor()
node = ast.parse('1/1')
value_obj = crep.cpp_value('i', gc_scope_top_level(), ctyp.terminal('int'))
sequence = crep.cpp_sequence(
value_obj,
crep.cpp_value('i', gc_scope_top_level(), ctyp.terminal('int')),
gc_scope_top_level())
node.rep = sequence # type: ignore
as_root = q.get_as_ROOT(node)
assert isinstance(as_root, rh.cpp_ttree_rep)
def make_sequence_from_collection(self, rep):
'''
Take a collection and produce a sequence. Eventually this should likely be some sort of
plug-in architecture. But for now, we will just assume everything looks like a vector. When
it comes time for a new type, this is where it should go.
'''
element_type = rep.cpp_type().element_type()
iterator_value = crep.cpp_value(unique_name("i_obj"), None, element_type) # type: ignore
l_statement = statement.loop(iterator_value, crep.dereference_var(rep)) # type: ignore
self._gc.add_statement(l_statement)
iterator_value.reset_scope(self._gc.current_scope())
# For a new sequence like this the sequence and iterator value are the same
return crep.cpp_sequence(iterator_value, iterator_value, self._gc.current_scope())
def test_as_root_as_dict():
q = query_ast_visitor()
node = ast.parse('1/1')
dict_obj = crep.cpp_dict(
{
ast.Constant(value='hi'):
crep.cpp_value('i', gc_scope_top_level(), ctyp.terminal('int'))
}, gc_scope_top_level())
sequence = crep.cpp_sequence(
dict_obj, # type: ignore
crep.cpp_value('i', gc_scope_top_level(), ctyp.terminal('int')),
gc_scope_top_level())
node.rep = sequence # type: ignore
as_root = q.get_as_ROOT(node)
assert isinstance(as_root, rh.cpp_ttree_rep)
async def exe_from_qastle(q: str):
'Dummy executor that will return the ast properly rendered. If qastle_roundtrip is true, then we will round trip the ast via qastle first.'
# Round trip qastle if requested.
import qastle
a = qastle.text_ast_to_python_ast(q).body[0].value
# Setup the rep for this filter
from func_adl import find_EventDataset
file = find_EventDataset(a)
iterator = cpp_variable("bogus-do-not-use",
top_level_scope(),
cpp_type=None)
file.rep = cpp_sequence(iterator, iterator,
top_level_scope()) # type: ignore
# Use the dummy executor to process this, and return it.
exe = dummy_executor()
exe.evaluate(a)
return exe
def call_Select(self, node: ast.Call, args: List[ast.arg]):
'Transform the iterable from one form to another'
assert len(args) == 2
source = args[0]
selection = cast(ast.Lambda, args[1])
# Make sure we are in a loop
seq = self.as_sequence(source)
# Simulate this as a "call"
selection = lambda_unwrap(selection)
c = ast.Call(func=selection, args=[seq.sequence_value().as_ast()])
new_sequence_value = self.get_rep(c)
# We need to build a new sequence.
rep = crep.cpp_sequence(new_sequence_value, seq.iterator_value(), self._gc.current_scope())
node.rep = rep # type: ignore
self._result = rep
return rep
async def execute_result_async(self, a: ast.AST) -> Any:
'Dummy executor that will return the ast properly rendered. If qastle_roundtrip is true, then we will round trip the ast via qastle first.'
# Round trip qastle if requested.
if self._q_roundtrip:
import qastle
print(f'before: {ast.dump(a)}')
a_text = qastle.python_ast_to_text_ast(a)
a = qastle.text_ast_to_python_ast(a_text).body[0].value
print(f'after: {ast.dump(a)}')
# Setup the rep for this dataset
from func_adl import find_EventDataset
file = find_EventDataset(a)
iterator = cpp_variable("bogus-do-not-use",
top_level_scope(),
cpp_type=None)
file.rep = cpp_sequence(iterator, iterator,
top_level_scope()) # type: ignore
# Use the dummy executor to process this, and return it.
exe = dummy_executor()
rnr = atlas_xaod_executor()
exe.evaluate(a)
return exe
def get_as_ROOT(self, node: ast.AST) -> rh.cpp_ttree_rep:
'''For a given node, return a root ttree rep.
This is used to make sure whatever the sequence is, that it returns a RootTTree. Use this
when the user may not have explicitly requested an output format. For example, if they end
with a tuple or a dict request, but not a AsAwkwardArray or similar.
1. If the top level ast node already requests an ROOTTTree, the representation is returned.
1. If the node is a sequence of `cpp_values`, then a single column tree is created.
1. If the node is a sequence of dict's, a ttree is created that uses the dictionary
key's as column names.
1. Anything else causes an exception to be raised.
Args:
node (ast.AST): Top level `func_adl` expression that is to be renered as a ROOT file.
Returns:
rh.cpp_ttree_rep: The resulting node that will generate a root file.
Exceptions:
ValueError: If we end with something other than above, we raise `ValueError` to
indicate that we can't figure out how to convert something into a ROOT
file.
'''
r = self.get_rep(node)
if isinstance(r, rh.cpp_ttree_rep):
return r
# If this isn't a sequence, then we are totally blown here.
if not isinstance(r, crep.cpp_sequence):
raise ValueError(f'Do not know how to convert {r} into a ROOT file')
# If this is a dict, then pull out each item and re-assemble into a tuple
# which we can feed to the root guy.
values = r.sequence_value()
if isinstance(values, crep.cpp_dict):
values = cast(crep.cpp_dict, values)
col_values = values.value_dict.values()
col_names = ast.List(elts=list(values.value_dict.keys()))
s_tuple = crep.cpp_tuple(tuple(col_values), values.scope())
tuple_sequence = crep.cpp_sequence(s_tuple, r.iterator_value(), r.scope()) # type: ignore
ast_dummy_source = ast.Constant(value=1)
ast_dummy_source.rep = tuple_sequence # type: ignore
ast_ttree = function_call('ResultTTree',
[ast_dummy_source,
col_names,
ast.parse('"xaod_tree"').body[0].value, # type: ignore
ast.parse('"xaod_output"').body[0].value]) # type: ignore
result = self.get_rep(ast_ttree)
assert isinstance(result, rh.cpp_ttree_rep)
return result
if isinstance(values, crep.cpp_tuple):
col_names = ast.List(elts=[ast.parse(f"'col{i}'").body[0].value for i, _ in enumerate(values.values())]) # type: ignore
ast_ttree = function_call('ResultTTree',
[node,
col_names,
ast.parse('"xaod_tree"').body[0].value, # type: ignore
ast.parse('"xaod_output"').body[0].value]) # type: ignore
result = self.get_rep(ast_ttree)
assert isinstance(result, rh.cpp_ttree_rep)
return result
elif isinstance(values, crep.cpp_value) or isinstance(values, crep.cpp_sequence):
ast_ttree = function_call('ResultTTree',
[node,
ast.parse('"col1"').body[0].value, # type: ignore
ast.parse('"xaod_tree"').body[0].value, # type: ignore
ast.parse('"xaod_output"').body[0].value]) # type: ignore
result = self.get_rep(ast_ttree)
assert isinstance(result, rh.cpp_ttree_rep)
return result
else:
raise ValueError(f'Do not know how to convert a sequence of {r.sequence_value()} into a ROOT file.')
raise NotImplementedError()
