def test_eds_recovery_two_ds():
r1 = my_event()
r2 = my_event()
q1 = r1.Select(lambda a: a + 1)
q2 = r2.Select(lambda b: b + 1)
q = ObjectStream(ast.BinOp(q1.query_ast, ast.Add, q2.query_ast))
with pytest.raises(Exception) as e:
find_EventDataset(q.query_ast)
assert "more than one" in str(e)
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'
])
async def _get_query(self) -> str:
"""Return the qastle query.
Note: To do this we have to forward-cast the object: by design, not all `func_adl`
queries are `ServiceX` queries. But this library only works with datasets that are
based in `ServiceX`. Thus some duck typing occurs in this method.
"""
event_dataset_ast = find_EventDataset(self.query.query_ast)
event_dataset = event_dataset_ast._eds_object # type: ignore
if not hasattr(event_dataset, "return_qastle"):
raise Exception(
f"Base func_adl query {str(event_dataset)} does not have a way to generate qastle!"
)
event_dataset.return_qastle = True # type: ignore
return await self.query.value_async()
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)
set_rep(file, cpp_sequence(iterator, iterator, top_level_scope()))
# Use the dummy executor to process this, and return it.
exe = atlas_xaod_dummy_executor()
exe.evaluate(a)
return exe
async def execute_result_async(self, a: ast.AST, title: str) -> 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)
set_rep(file, cpp_sequence(iterator, iterator, top_level_scope()))
# Use the dummy executor to process this, and return it.
exe = self.get_dummy_executor_obj()
exe.evaluate(a)
return exe
def write_cpp_files(self, ast: ast.AST,
output_path: Path) -> ExecutionInfo:
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)
crep.set_rep(file,
crep.cpp_sequence(iterator, iterator, top_level_scope()))
# Visit the AST to generate the code structure and find out what the
# result is going to be.
qv = self.get_visitor_obj()
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_decl_code = qv.class_declaration_code()
includes = qv.include_files() + self.body_include_files
link_libraries = qv.link_libraries() + self.link_libraries
# The replacement dict to pass to the template generator can now be filled
info = {}
info['query_code'] = query_code.lines_of_query_code()
info['class_decl'] = class_decl_code
info['book_code'] = book_code.lines_of_query_code()
info['body_include_files'] = includes
info['header_include_files'] = self.header_include_files
info['private_members'] = self.private_members
info['instance_initialization'] = self.instance_initialization
info['initialize_lines'] = self.initialize_lines
info['ctor_lines'] = self.ctor_lines
info['link_libraries'] = link_libraries
info.update(self.add_to_replacement_dict())
# We use jinja2 templates. Write out everything.
template_dir = _find_dir(self._template_dir_name)
j2_env = jinja2.Environment(
loader=jinja2.FileSystemLoader(template_dir))
for file_name in self._file_names:
self._copy_template_file(j2_env, info, file_name, output_path)
(output_path / self._runner_name).chmod(0o755)
# Reset our object for the next call (e.g. reset global state)
self.reset()
# Build the return object.
return ExecutionInfo(result_rep, output_path, self._runner_name,
self._file_names)
def test_eds_recovery_with_odd_call():
r = my_event()
q = r.Select(lambda a: (lambda b: b + 1)(a))
found_node = find_EventDataset(q.query_ast)
assert found_node._eds_object == r # type: ignore
def test_eds_recovery_no_root():
q = ObjectStream(ast.BinOp(ast.Num(1), ast.Add, ast.Num(2)))
with pytest.raises(Exception) as e:
find_EventDataset(q.query_ast)
assert "no root" in str(e)
def test_eds_recovery():
"Make sure we can get back the event dataset"
r = my_event()
found_node = find_EventDataset(r.query_ast)
assert hasattr(found_node, "_eds_object")
assert found_node._eds_object == r # type: ignore
