def apply_ast_transformations(self, a: ast.AST):
r'''
Run through all the transformations that we have on tap to be run on the client side.
Return a (possibly) modified ast.
'''
# Do tuple resolutions. This might eliminate a whole bunch fo code!
a = change_extension_functions_to_calls(a)
a = aggregate_node_transformer().visit(a)
a = simplify_chained_calls().visit(a)
a = find_known_functions().visit(a)
# Any C++ custom code needs to be threaded into the ast
a = cpp_ast.cpp_ast_finder().visit(a)
# And return the modified ast
return a
def apply_ast_transformations(self, a: ast.AST):
r'''
Run through all the transformations that we have on tap to be run on the client side.
Return a (possibly) modified ast.
'''
# Do tuple resolutions. This might eliminate a whole bunch fo code!
a, meta_data = extract_metadata(a)
cpp_functions = process_metadata(meta_data)
a = change_extension_functions_to_calls(a)
a = aggregate_node_transformer().visit(a)
a = simplify_chained_calls().visit(a)
a = find_known_functions().visit(a)
# Any C++ custom code needs to be threaded into the ast
method_names = dict(self._method_names)
method_names.update({
md.name: (lambda call_node, md=md: cpp_ast.build_CPPCodeValue(
md, call_node))
if isinstance(md, cpp_ast.CPPCodeSpecification) # type: ignore
else self.build_collection_callback(md)
for md in cpp_functions
if isinstance(md, (cpp_ast.CPPCodeSpecification,
EventCollectionSpecification))
})
a = cpp_ast.cpp_ast_finder(method_names).visit(a)
# Save the injection blocks
self._inject_blocks = [
md for md in cpp_functions if isinstance(md, InjectCodeBlock)
]
# Pull off any joboption blocks
for m in cpp_functions:
if isinstance(m, JobScriptSpecification):
self._job_option_blocks.append(m)
# And return the modified ast
return a
def test_extension_functions_select_extension_in_lambda_too():
source = ast.parse("jets.Select(lambda b: b.Select(lambda j: jpt()))")
expected = ast.parse("Select(jets, lambda b: Select(b, lambda j: jpt()))")
transform = change_extension_functions_to_calls(source)
assert ast.dump(transform) == ast.dump(expected)
def test_extension_functions_call():
source = ast.parse("dude()")
expected = ast.parse("dude()")
transform = change_extension_functions_to_calls(source)
assert ast.dump(transform) == ast.dump(expected)
