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