def visit_Call_Aggregate_only(self, node: ast.Call):
        '''
        - (acc lambda): the accumulator is set to the first element, and the lambda is called to
                        update it after that. This is called `agg_only`.
        '''
        agg_lambda = node.args[0]

        # Get the sequence we are calling against and the accumulator
        if not isinstance(node.func, ast.Attribute):
            raise BaseException("Wrong type of function")
        seq = self.as_sequence(node.func.value)
        accumulator, accumulator_scope = self.create_accumulator(seq)

        # We have to do a simple if statement here so that the first time through we can set the
        # accumulator, and the second time we can add to it.

        is_first_iter = crep.cpp_variable(unique_name("is_first"),
                                          self._gc.current_scope(),
                                          cpp_type=ctyp.terminal('bool'),
                                          initial_value=crep.cpp_value(
                                              'true', self._gc.current_scope(),
                                              ctyp.terminal('bool')))
        accumulator_scope.declare_variable(is_first_iter)

        # Set the scope where we will be doing the accumulation
        sv = seq.sequence_value()
        if isinstance(sv, crep.cpp_sequence):
            self._gc.set_scope(sv.iterator_value().scope()[-1])
        else:
            self._gc.set_scope(sv.scope())

        # Code up if statement to select out the first element.
        if_first = statement.iftest(is_first_iter)
        self._gc.add_statement(if_first)
        self._gc.add_statement(
            statement.set_var(
                is_first_iter,
                crep.cpp_value("false", self._gc.current_scope(),
                               ctyp.terminal('bool'))))

        # Set the accumulator
        self._gc.add_statement(
            statement.set_var(accumulator, seq.sequence_value()))
        self._gc.pop_scope()

        # Now do the if statement and make the call to calculate the accumulation.
        self._gc.add_statement(statement.elsephrase())
        call = ast.Call(
            func=agg_lambda,
            args=[accumulator.as_ast(),
                  seq.sequence_value().as_ast()])
        self._gc.add_statement(
            statement.set_var(accumulator, self.get_rep(call)))

        # Finally, since this is a terminal, we need to pop off the top.
        self._gc.set_scope(accumulator_scope)

        # Cache the results in our result in case we are skipping nodes in the AST.
        node.rep = accumulator
        self._result = accumulator
Example #2
0
def test_expression_pointer_decl():
    e2 = crep.cpp_value("dude", top_level_scope(), ctyp.terminal("int"))
    assert False == e2.is_pointer()

    e3 = crep.cpp_value("dude", top_level_scope(),
                        ctyp.terminal("int", is_pointer=True))
    assert True == e3.is_pointer()
def test_deepest_scope_equal():
    g = generated_code()
    s1 = statement.iftest("true")
    s2 = 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)
Example #4
0
def getAttributeFloatAst(call_node: ast.Call):
    r'''
    Return an attribute on one of the xAOD objects.
    '''
    # Get the name of the moment out
    if len(call_node.args) != 1:
        raise BaseException(
            "Calling getMomentFloat - only one argument is allowed")
    if not isinstance(call_node.args[0], ast.Str):
        raise BaseException(
            "Calling getMomentFloat - only acceptable argument is a string")

    r = cpp_ast.CPPCodeValue()
    # We don't need include files for this - just quick access
    r.args = [
        'moment_name',
    ]
    r.replacement_instance_obj = ('obj_j', call_node.func.value.id)
    r.running_code += [
        'float result = obj_j->getAttribute<float>(moment_name);'
    ]
    r.result = 'result'
    r.result_rep = lambda sc: crep.cpp_variable(
        unique_name("jet_attrib"), scope=sc, cpp_type=ctyp.terminal('float'))

    # Replace it as the function that is going to get called.
    call_node.func = r

    return call_node
def determine_type_mf(parent_type, function_name):
    '''
    Determine the return type of the member function. Do our best to make
    an intelligent case when we can.

    parent_type:        the type of the parent
    function_name:      the name of the function we are calling
    '''
    # If we don't know the type...
    if parent_type is None:
        raise BaseException(
            "Internal Error: Trying to call member function for a type we do not know!"
        )
    # If we are doing one of the normal "terminals", then we can just bomb. This should not happen!

    rtn_type = ctyp.method_type_info(str(parent_type), function_name)
    if rtn_type is not None:
        return rtn_type

    # We didn't know it. Lets make a guess, and error out if we are clearly making a mistake.
    base_types = ['double', 'float', 'int']
    s_parent_type = str(parent_type)
    if s_parent_type in base_types:
        raise xAODTranslationError(
            f'Unable to call method {function_name} on type {str(parent_type)}.'
        )

    # Ok - we give up. Return a double.
    print(
        f"Warning: assumping that the method '{str(s_parent_type)}.{function_name}(...)' has return type 'double'. Use cpp_types.add_method_type_info to suppress (or correct) this warning."
    )
    return ctyp.terminal('double')
Example #6
0
def getAttributeVectorFloatAst(call_node: ast.Call):
    r'''
    Return a cpp ast accessing a vector of doubles for an xAOD attribute
    '''
    # Get the name of the moment out
    if len(call_node.args) != 1:
        raise BaseException(
            "Calling getMomentFloat - only one argument is allowed")
    if not isinstance(call_node.args[0], ast.Str):
        raise BaseException(
            "Calling getMomentFloat - only acceptable argument is a string")

    r = cpp_ast.CPPCodeValue()
    r.include_files += ['vector']
    r.args = [
        'moment_name',
    ]
    r.replacement_instance_obj = ('obj_j', call_node.func.value.id)
    r.running_code += [
        'auto result = obj_j->getAttribute<std::vector<double>>(moment_name);'
    ]
    r.result = 'result'
    r.result_rep = lambda sc: crep.cpp_collection(
        unique_name("jet_vec_attrib_"),
        scope=sc,
        collection_type=ctyp.collection(ctyp.terminal('double')))

    # Replace it as the function that is going to get called.
    call_node.func = r

    return call_node
    def visit_IfExp(self, node):
        r'''
        We'd like to be able to use the "?" operator in C++, but the
        problem is lazy evaluation. It could be when we look at one or the
        other item, a bunch of prep work has to be done - and that will
        show up in separate statements. So we have to use if/then/else with
        a result value.
        '''

        # The result we'll store everything in.
        result = crep.cpp_variable(unique_name("if_else_result"),
                                   self._gc.current_scope(),
                                   cpp_type=ctyp.terminal("double"))
        self._gc.declare_variable(result)

        # We always have to evaluate the test.
        current_scope = self._gc.current_scope()
        test_expr = self.get_rep(node.test)
        self._gc.add_statement(statement.iftest(test_expr))
        if_scope = self._gc.current_scope()

        # Next, we do the true and false if statement.
        self._gc.add_statement(
            statement.set_var(result, self.get_rep(node.body)))
        self._gc.set_scope(if_scope)
        self._gc.pop_scope()
        self._gc.add_statement(statement.elsephrase())
        self._gc.add_statement(
            statement.set_var(result, self.get_rep(node.orelse)))
        self._gc.set_scope(current_scope)

        # Done, the result is the rep of this node!
        node.rep = result
        self._result = result
    def visit_First(self, node):
        'We are in a sequence. Take the first element of the sequence and use that for future things.'

        # Make sure we are in a loop.
        seq = self.as_sequence(node.source)

        # The First terminal works by protecting the code with a if (first_time) {} block.
        # We need to declare the first_time variable outside the block where the thing we are
        # looping over here is defined. This is a little tricky, so we delegate to another method.
        loop_scope = seq.iterator_value().scope()
        outside_block_scope = loop_scope[-1]

        # Define the variable to track this outside that block.
        is_first = crep.cpp_variable(unique_name('is_first'),
                                     outside_block_scope,
                                     cpp_type=ctyp.terminal('bool'),
                                     initial_value=crep.cpp_value(
                                         'true', self._gc.current_scope(),
                                         ctyp.terminal('bool')))
        outside_block_scope.declare_variable(is_first)

        # Now, as long as is_first is true, we can execute things inside this statement.
        # The trick is putting the if statement in the right place. We need to locate it just one level
        # below where we defined the scope above.
        s = statement.iftest(is_first)
        s.add_statement(
            statement.set_var(
                is_first,
                crep.cpp_value('false',
                               top_level_scope(),
                               cpp_type=ctyp.terminal('bool'))))

        sv = seq.sequence_value()
        if isinstance(sv, crep.cpp_sequence):
            self._gc.set_scope(sv.iterator_value().scope()[-1])
        else:
            self._gc.set_scope(sv.scope())
        self._gc.add_statement(s)

        # If we just found the first sequence in a sequence, return that.
        # Otherwise return a new version of the value.
        first_value = sv if isinstance(
            sv, crep.cpp_sequence) else sv.copy_with_new_scope(
                self._gc.current_scope())

        node.rep = first_value
        self._result = first_value
def test_can_call_prodVtx():
    ctyp.add_method_type_info(
        "xAOD::TruthParticle", "prodVtx",
        ctyp.terminal('xAODTruth::TruthVertex', is_pointer=True))
    EventDataset("file://root.root") \
        .Select("lambda e: e.TruthParticles('TruthParticles').Select(lambda t: t.prodVtx().x()).Sum()") \
        .AsROOTTTree('root.root', 'dude', "n_jets") \
        .value(executor=exe_for_test)
    def visit_Compare(self, node):
        'A compare between two things. Python supports more than that, but not implemented yet.'
        if len(node.ops) != 1:
            raise BaseException("Do not support 1 < a < 10 comparisons yet!")

        left = self.get_rep(node.left)
        right = self.get_rep(node.comparators[0])

        r = crep.cpp_value(
            '({0}{1}{2})'.format(left.as_cpp(),
                                 compare_operations[type(node.ops[0])],
                                 right.as_cpp()), self._gc.current_scope(),
            ctyp.terminal("bool"))
        node.rep = r
        self._result = r
def test_int_pointer():
    t_int = ctyp.terminal('int')
    assert False == t_int.is_pointer()
 def __init__(self, filename, treename, scope):
     cpp_value.__init__(self, unique_name("pandas"), scope,
                        ctyp.terminal("pandasdf"))
     self.filename = filename
     self.treename = treename
 def __init__(self, filename, treename, scope):
     cpp_value.__init__(self, unique_name("awk_array"), scope,
                        ctyp.terminal("awkwardarray"))
     self.filename = filename
     self.treename = treename
 def __init__(self, filename, treename, scope):
     cpp_value.__init__(self, unique_name("ttree_rep"), scope,
                        ctyp.terminal("ttreetfile"))
     self.filename = filename
     self.treename = treename
Example #15
0
 def element_type(self):
     'Return the type of the elements in the collection'
     return ctyp.terminal(self._element_name, is_pointer=True)
Example #16
0
    r.running_code += ['{0} result = 0;'.format(info['container_type']),
                       'ANA_CHECK (evtStore()->retrieve(result, collection_name));']
    r.result = 'result'

    is_collection = info['is_collection'] if 'is_collection' in info else True
    if is_collection:
        r.result_rep = lambda scope: crep.cpp_collection(unique_name(info['function_name'].lower()), scope=scope, collection_type=info['container_type'])
    else:
        r.result_rep = lambda scope: crep.cpp_variable(unique_name(info['function_name'].lower()), scope=scope, cpp_type=info['container_type'])

    # Replace it as the function that is going to get called.
    call_node.func = r

    return call_node


# Config everything.
def create_higher_order_function(info):
    'Creates a higher-order function because python scoping is broken'
    return lambda call_node: getCollection(info, call_node)


for info in collections:
    cpp_ast.method_names[info['function_name']] = create_higher_order_function(info)


# Configure some info about the types.
# TODO: Make a way to do this in client programs, or automate the parsing of types
ctyp.add_method_type_info("xAOD::TruthParticle", "prodVtx", ctyp.terminal('xAODTruth::TruthVertex', is_pointer=True))
ctyp.add_method_type_info("xAOD::TruthParticle", "decayVtx", ctyp.terminal('xAODTruth::TruthVertex', is_pointer=True))
def guess_type_from_number(n):
    if int(n) == n:
        return ctyp.terminal("int")
    return ctyp.terminal("double")
def test_method_type_found():
    ctyp.add_method_type_info("bogus", "pt", ctyp.terminal('double'))
    assert 'double' == str(ctyp.method_type_info("bogus", "pt"))
 def visit_Str(self, node):
     node.rep = crep.cpp_value('"{0}"'.format(node.s),
                               self._gc.current_scope(),
                               ctyp.terminal("string"))
     self._result = node.rep