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)