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)
assert seq.sequence_value().cpp_type().type == 'int'
def test_variable_type_update():
tc = gc_scope_top_level()
expr = "a"
ctype = ctyp.terminal('int', False)
v = crep.cpp_variable(expr, tc, ctype)
v.update_type(ctyp.terminal('float', False))
assert v.cpp_type().type == 'float'
def define_default_atlas_types():
'Define the default atlas types'
ctyp.add_method_type_info(
"xAOD::TruthParticle", "prodVtx",
ctyp.terminal('xAODTruth::TruthVertex', p_depth=1))
ctyp.add_method_type_info(
"xAOD::TruthParticle", "decayVtx",
ctyp.terminal('xAODTruth::TruthVertex', p_depth=1))
ctyp.add_method_type_info("xAOD::TruthParticle", "parent",
ctyp.terminal('xAOD::TruthParticle', p_depth=1))
ctyp.add_method_type_info("xAOD::TruthParticle", "child",
ctyp.terminal('xAOD::TruthParticle', p_depth=1))
def test_as_root_as_single_column():
q = atlas_xaod_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 test_deepest_scope_equal():
g = generated_code()
s1 = statement.iftest("true")
statement.set_var("v1", "true")
g.add_statement(s1)
scope_1 = g.current_scope()
v1 = crep.cpp_value("v1", scope_1, ctyp.terminal('int'))
v2 = crep.cpp_value("v2", scope_1, ctyp.terminal('int'))
assert v1 == deepest_scope(v1, v2)
assert v2 == deepest_scope(v2, v1)
def test_variable_type__with_initial_update():
tc = gc_scope_top_level()
expr = "a"
c_type = ctyp.terminal('int', False)
c_init = crep.cpp_value('0.0', tc, ctyp.terminal('int', False))
v = crep.cpp_variable(expr, tc, c_type, c_init)
v.update_type(ctyp.terminal('float', False))
assert v.cpp_type().type == 'float'
iv = v.initial_value()
assert iv is not None
assert iv.cpp_type().type == 'float'
def test_deepest_scope_one_greater():
g = generated_code()
s1 = statement.iftest("true")
s2 = statement.iftest("true")
g.add_statement(s1)
scope_1 = g.current_scope()
g.add_statement(s2)
scope_2 = g.current_scope()
v1 = crep.cpp_value("v1", scope_1, ctyp.terminal('int'))
v2 = crep.cpp_value("v2", scope_2, ctyp.terminal('int'))
assert v2 == deepest_scope(v1, v2)
assert v2 == deepest_scope(v2, v1)
def test_can_call_prodVtx():
ctyp.add_method_type_info(
"xAOD::TruthParticle", "prodVtx",
ctyp.terminal('xAODTruth::TruthVertex', p_depth=1))
atlas_xaod_dataset("file://root.root") \
.Select("lambda e: e.TruthParticles('TruthParticles').Select(lambda t: t.prodVtx().x()).Sum()") \
.value()
def test_as_root_as_dict():
q = atlas_xaod_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)
def test_cpp_value_as_str():
'Make sure we can generate a str from a value - this will be important for errors'
v1 = crep.cpp_value('dude', top_level_scope(), ctyp.terminal('int'))
assert 'dude' in str(v1)
v2 = crep.cpp_value('dude', top_level_scope(), None)
assert 'dude' in str(v2)
def test_variable_pointer_2():
'Make sure p_depth can deal with a non-type correctly'
v1 = crep.cpp_value('dude', top_level_scope(), ctyp.terminal('int'))
v2 = crep.cpp_value('dude', top_level_scope(), None)
assert v1.cpp_type().type == 'int'
with pytest.raises(RuntimeError):
v2.cpp_type()
def test_collection_return():
'Make sure we can set and deal with a returned collection properly'
ctyp.add_method_type_info(
"xAOD::TruthParticle", "vertexes",
ctyp.collection(ctyp.terminal('double'), p_depth=1))
(atlas_xaod_dataset("file://root.root").SelectMany(
lambda e: e.TruthParticles('TruthParticles')).SelectMany(
lambda t: t.vertexes()).value())
def test_deref_simple_ptr():
tc = gc_scope_top_level()
expr = "a"
c_type = ctyp.terminal('int', 1)
v = crep.cpp_variable(expr, tc, c_type)
d = crep.dereference_var(v)
assert d.cpp_type().type == 'int'
assert d.cpp_type().p_depth == 0
assert d.as_cpp() == '*a'
def test_compare_string_var():
q = atlas_xaod_query_ast_visitor()
node = ast.parse('e == "hi"').body[0].value # type: ignore
node.left.rep = crep.cpp_value('e', gc_scope_top_level(),
ctyp.terminal('string')) # type: ignore
r = q.get_rep(node)
assert isinstance(r, crep.cpp_value)
assert r.cpp_type().type == 'bool'
assert r.as_cpp() == '(e=="hi")'
def add_function_mapping(python_name, cpp_name, include_files, return_type):
'''Add a re-mapping from a python function to an actual function.
python_name: fully qualified name of the python function
cpp_name: fully qualified name of the C++ function
include_files: any include files that should be included in the C++ source. Can also be a single string.
return_type: C++ return type
'''
global functions_to_replace
functions_to_replace[python_name] = cpp_function(
cpp_name, include_files if type(include_files) is list else [
include_files,
], terminal(return_type))
def test_deref_collection_ptr():
tc = gc_scope_top_level()
c_type = ctyp.collection(ctyp.terminal(ctyp.parse_type('int')), ctyp.parse_type('vector<int>*'))
c = crep.cpp_collection('my_var', tc, c_type)
d = crep.dereference_var(c)
assert isinstance(d, crep.cpp_collection)
cpp_type = d.cpp_type()
assert isinstance(cpp_type, ctyp.collection)
assert str(cpp_type) == 'vector<int>'
assert str(cpp_type.element_type) == 'int'
def test_subscript():
q = atlas_xaod_query_ast_visitor()
our_a = ast.Name(id='a')
our_a.rep = crep.cpp_collection('jets', gc_scope_top_level(),
ctyp.collection(
ctyp.terminal('int'))) # type: ignore
node = ast_parse_with_replacement('a[10]', {
'a': our_a
}).body[0].value # type: ignore
as_root = q.get_rep(node)
assert isinstance(as_root, crep.cpp_value)
assert str(as_root.cpp_type()) == 'int'
assert as_root.as_cpp() == 'jets.at(10)'
def test_member_access_obj_ptr():
cv = crep.cpp_value('f', gc_scope_top_level(), ctyp.terminal(ctyp.parse_type('obj*')))
assert crep.base_type_member_access(cv) == 'f->'
def __init__(self, filename, treename, scope):
cpp_value.__init__(self, unique_name("ttree_rep"), scope, ctyp.terminal("ttreetfile"))
self.filename = filename
self.treename = treename
def test_expression_pointer_decl():
e2 = crep.cpp_value("dude", top_level_scope(), ctyp.terminal("int"))
assert e2.p_depth == 0
e3 = crep.cpp_value("dude", top_level_scope(), ctyp.terminal("int", p_depth=1))
assert e3.p_depth == 1
def test_member_access_obj_depth_1():
cv = crep.cpp_value('f', gc_scope_top_level(), ctyp.terminal(ctyp.parse_type('obj')))
assert crep.base_type_member_access(cv, 2) == '(*f)->'
def test_accurate_type_float_and_double():
t1 = ctyp.terminal('double', False)
t2 = ctyp.terminal('float', False)
r = most_accurate_type([t1, t2])
assert r._type == 'double'
def test_accurate_type_single_int():
t = ctyp.terminal('int', False)
r = most_accurate_type([t])
assert r._type == 'int'
def test_accurate_type_int_and_float():
t1 = ctyp.terminal('int', False)
t2 = ctyp.terminal('float', False)
r = most_accurate_type([t1, t2])
assert r._type == 'float'
def test_accurate_type_two_int():
t1 = ctyp.terminal('int', False)
t2 = ctyp.terminal('int', False)
r = most_accurate_type([t1, t2])
assert r._type == 'int'
def isNonnullAst(call_node):
r'''
User is trying to test, on certian objects, if the object, when dereferenced, will be a null
pointer. This is tricky for our data model because it treats an object as both a pointer and
a object. Most of `func_adl_xAOD` is setup to deal with it as a pointer. This allows us to
get around ths.
'''
if len(call_node.args) != 1:
raise ValueError("Calling isNonnull(object) has incorrect number of arguments")
# Create an AST to hold onto all of this.
r = cpp_ast.CPPCodeValue()
# We need all four arguments pushed through.
r.args = ['cms_object']
# The code is three steps
r.running_code += ['auto result = (cms_object).isNonnull();']
r.result = 'result'
r.result_rep = lambda scope: cpp_variable(unique_name('is_non_null'), scope=scope, cpp_type=ctyp.terminal('bool'))
call_node.func = r
return call_node
def build_CPPCodeValue(spec: CPPCodeSpecification,
call_node: ast.Call) -> ast.Call:
'''
Given the specification for a C++ code block, invoked as a function in our AST, replace
the call node with a cpp spec callback AST so the C++ code is properly inserted into the
call stream.
Args:
spec (CPPCodeSpecification): The specification, including the code, that we should insert at this call node
call_node (ast.Call): The call node (with arguments) that we are going to replace with this
C++ code.
Raises:
ValueError: Raised if something is wrong with the call site
Returns:
[type]: The C++ ast that can easily be emitted as code
'''
if len(call_node.args) != len(spec.arguments):
raise ValueError(
f"The call of {spec.name}({', '.join(spec.arguments)}) has insufficient arguments ({len(call_node.args)})."
)
if isinstance(call_node.func,
ast.Attribute) and spec.method_object is None:
raise ValueError(
f"The {spec.name} is a function, but it is invoked like a method.")
if isinstance(call_node.func, ast.Name) and spec.method_object is not None:
raise ValueError(
f"The {spec.name} is a method, but it is invoked like a function.")
# Create an AST to hold onto all of this.
r = CPPCodeValue()
# We need TVector2 included here
r.include_files += spec.include_files
# We need all four arguments pushed through.
r.args = spec.arguments
# The code is three steps
r.running_code += spec.code
r.result = spec.result
if spec.cpp_return_is_collection:
r.result_rep = lambda scope: cpp_collection(
unique_name(spec.name),
scope=scope,
collection_type=ctyp.collection(ctyp.terminal(spec.cpp_return_type)
)) # type: ignore
else:
r.result_rep = lambda scope: cpp_variable(unique_name(spec.name),
scope=scope,
cpp_type=ctyp.terminal(
spec.cpp_return_type))
# If this is a mehtod, copy the info over to generate the obj reference.
if spec.method_object is not None:
r.replacement_instance_obj = (spec.method_object,
call_node.func.value.id) # type: ignore
call_node.func = r # type: ignore
return call_node
def process_metadata(md_list: List[Dict[str, Any]]) -> List[SpecificationTypes]:
'''Process a list of metadata, in order.
Args:
md (List[Dict[str, str]]): The metadata to process
Returns:
List[X]: Metadata we've found
'''
cpp_funcs: List[SpecificationTypes] = []
for md in md_list:
md_type = md.get('metadata_type')
if md_type is None:
raise ValueError(f'Metadata is missing `metadata_type` info ({md})')
if md_type == 'add_method_type_info':
if 'return_type' in md:
# Single return type
type_info = parse_type(md['return_type'])
term = terminal(type_info.name, p_depth=type_info.pointer_depth)
else:
type_info_element = parse_type(md['return_type_element'])
type_info_collection = parse_type(md['return_type_collection']) if 'return_type_collection' in md else CPPParsedTypeInfo(f'std::vector<{type_info_element}>', 0)
term = collection(terminal(type_info_element), array_type=type_info_collection)
d_count = 0
if 'deref_count' in md:
d_count = int(md['deref_count'])
add_method_type_info(md['type_string'], md['method_name'], term, d_count)
elif md_type == 'inject_code':
info = dict(md)
del info['metadata_type']
if len(info) > 0:
try:
spec = InjectCodeBlock(**info)
except TypeError as e:
raise ValueError(f'Bad inject_code block item: {str(e)}')
if ok_to_add_code_block(spec, cpp_funcs):
cpp_funcs.append(spec)
elif md_type == 'add_job_script':
spec = JobScriptSpecification(
name=md['name'],
script=md['script'],
depends_on=md.get('depends_on', [])
)
cpp_funcs.append(spec)
elif md_type == 'add_cpp_function':
spec = CPPCodeSpecification(
md['name'],
md['include_files'],
md['arguments'],
md['code'],
md['result_name'] if 'result_name' in md else 'result',
md['return_type'],
bool(md['return_is_collection']) if 'return_is_collection' in md else False,
md['method_object'] if 'method_object' in md else None,
md['instance_object'] if 'instance_object' in md else None,
)
cpp_funcs.append(spec)
elif md_type == 'add_atlas_event_collection_info':
for k in md.keys():
if k not in ['metadata_type', 'name', 'include_files', 'container_type', 'element_type', 'contains_collection', 'link_libraries']:
raise ValueError(f'Unexpected key {k} when declaring ATLAS collection metadata')
if (md['contains_collection'] and 'element_type' not in md) or (not md['contains_collection'] and 'element_type' in md):
raise ValueError('In collection metadata, `element_type` must be specified if `contains_collection` is true and not if it is false')
from func_adl_xAOD.atlas.xaod.event_collections import atlas_xaod_event_collection_collection, atlas_xaod_event_collection_container
container_type = atlas_xaod_event_collection_collection(md['container_type'], md['element_type']) if md['contains_collection'] \
else atlas_xaod_event_collection_container(md['container_type'])
link_libraries = [] if 'link_libraries' not in md else md['link_libraries']
spec = EventCollectionSpecification(
'atlas',
md['name'],
md['include_files'],
container_type,
link_libraries)
cpp_funcs.append(spec)
elif md_type == 'add_cms_event_collection_info':
for k in md.keys():
if k not in ['metadata_type', 'name', 'include_files', 'container_type', 'element_type', 'contains_collection', 'element_pointer']:
raise ValueError(f'Unexpected key {k} when declaring ATLAS collection metadata')
if (md['contains_collection'] and 'element_type' not in md) or (not md['contains_collection'] and 'element_type' in md):
raise ValueError('In collection metadata, `element_type` must be specified if `contains_collection` is true and not if it is false')
from func_adl_xAOD.cms.aod.event_collections import cms_aod_event_collection_collection
container_type = cms_aod_event_collection_collection(md['container_type'], md['element_type'])
# container_type = cms_aod_event_collection_collection(md['container_type'], md['element_type']) if md['contains_collection'] \
# else cms_aod_event_collection_container(md['container_type'])
spec = EventCollectionSpecification(
'cms',
md['name'],
md['include_files'],
container_type,
[])
cpp_funcs.append(spec)
else:
raise ValueError(f'Unknown metadata type ({md_type})')
return cpp_funcs
def define_default_cms_types():
'Define the default cms types'
ctyp.add_method_type_info("reco::Track", "hitPattern", ctyp.terminal('reco::HitPattern'))
ctyp.add_method_type_info("reco::Muon", "globalTrack", ctyp.terminal('reco::Track', p_depth=1))
ctyp.add_method_type_info("reco::Muon", "hitPattern", ctyp.terminal('reco::HitPattern'))
ctyp.add_method_type_info("reco::Muon", "isPFIsolationValid", ctyp.terminal('bool'))
ctyp.add_method_type_info("reco::Muon", "isPFMuon", ctyp.terminal('bool'))
ctyp.add_method_type_info("reco::Muon", "pfIsolationR04", ctyp.terminal('reco::MuonPFIsolation'))
ctyp.add_method_type_info("reco::GsfElectron", "gsfTrack", ctyp.terminal('reco::GsfTrack', p_depth=1))
ctyp.add_method_type_info("reco::GsfElectron", "isEB", ctyp.terminal('bool'))
ctyp.add_method_type_info("reco::GsfElectron", "isEE", ctyp.terminal('bool'))
ctyp.add_method_type_info("reco::GsfElectron", "passingPflowPreselection", ctyp.terminal('bool'))
ctyp.add_method_type_info("reco::GsfElectron", "superCluster", ctyp.terminal('reco::SuperClusterRef', p_depth=1))
ctyp.add_method_type_info("reco::GsfElectron", "pfIsolationVariables", ctyp.terminal('reco::GsfElectron::PflowIsolationVariables'))
ctyp.add_method_type_info("reco::GsfTrack", "trackerExpectedHitsInner", ctyp.terminal('reco::HitPattern')) # reco::HitPattern is the expected return type
def __init__(self, type_name: Union[str, ctyp.CPPParsedTypeInfo],
element_name: Union[str, ctyp.CPPParsedTypeInfo],
p_depth_type: int, p_depth_element: int):
super().__init__(ctyp.terminal(element_name, p_depth=p_depth_element),
array_type=type_name,
p_depth=p_depth_type)
