def visit_call_Aggregate_initial(self, node: ast.Call):
'''
- (const, acc lambda): the accumulator is set to the value, and then the lambda is called to
update it on every single element. This is called `agg_initial`
'''
agg_lambda = node.args[1]
init_val = self.get_rep(node.args[0])
# Get the sequence we are calling against and the accumulator
seq = self.as_sequence(node.func.value)
accumulator, accumulator_scope = self.create_accumulator(
seq, initial_value=init_val, acc_type=init_val.cpp_type())
# Now do the accumulation. This happens at the current iterator scope.
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())
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
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
def call_Select(self, node: ast.Call, args: List[ast.arg]):
'Transform the iterable from one form to another'
assert len(args) == 2
source = args[0]
selection = cast(ast.Lambda, args[1])
# Make sure we are in a loop
seq = self.as_sequence(source)
# Simulate this as a "call"
selection = lambda_unwrap(selection)
c = ast.Call(func=selection, args=[seq.sequence_value().as_ast()])
new_sequence_value = self.get_rep(c)
# We need to build a new sequence.
rep = crep.cpp_sequence(new_sequence_value, seq.iterator_value(), self._gc.current_scope())
node.rep = rep # type: ignore
self._result = rep
return rep
def visit_Call(self, call_node: ast.Call):
r'''
Very limited call forwarding.
'''
# What kind of a call is this?
if isinstance(call_node.func, ast.Lambda):
self.visit_Call_Lambda(call_node)
elif isinstance(call_node.func, ast.Attribute):
self.visit_Call_Member(call_node)
elif isinstance(call_node.func, cpp_ast.CPPCodeValue):
self._result = cpp_ast.process_ast_node(self, self._gc, call_node)
elif isinstance(call_node.func, FunctionAST):
self._result = self.visit_function_ast(call_node)
else:
# Perhaps a method call we can normalize?
r = FuncADLNodeVisitor.visit_Call(self, call_node)
if r is None and not hasattr(call_node, 'rep'):
raise Exception("Do not know how to call '{0}'".format(ast.dump(call_node.func, annotate_fields=False)))
if r is not None:
self._result = r
call_node.rep = self._result # type: ignore
def visit_call_Aggregate_initial(self, node: ast.Call, args: List[ast.AST]):
'''
- (const, acc lambda): the accumulator is set to the value, and then the lambda is called to
update it on every single element. This is called `agg_initial`
'''
raw_seq = node.args[0]
init_val = self.get_rep(node.args[1])
agg_lambda = node.args[2]
assert isinstance(agg_lambda, ast.Lambda)
# Get the sequence we are calling against and the accumulator
seq = self.as_sequence(raw_seq)
accumulator, accumulator_scope = self._create_accumulator(seq, initial_value=init_val, acc_type=init_val.cpp_type())
# Now do the accumulation. This happens at the current iterator scope.
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())
call = ast.Call(func=agg_lambda, args=[accumulator.as_ast(), seq.sequence_value().as_ast()])
update_lambda = self.get_rep(call)
# Check the accumulator value still hols out. Since we need the accumulator previously,
# this will allow us to patch things up. This isn't perfect, but it will do.
if update_lambda.cpp_type().type != init_val.cpp_type().type:
best_type = most_accurate_type([init_val.cpp_type(), update_lambda.cpp_type()])
accumulator.update_type(best_type)
self._gc.add_statement(statement.set_var(accumulator, update_lambda))
# 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 # type: ignore
self._result = accumulator
def call_ResultTTree(self, node: ast.Call, args: List[ast.AST]):
'''This AST means we are taking an iterable and converting it to a ROOT file.
'''
# Unpack the variables.
assert len(args) == 4
source = args[0]
column_names = _extract_column_names(args[1])
tree_name = ast.literal_eval(args[2])
assert isinstance(tree_name, str)
# root_filename = args[3]
# Get the representations for each variable. We expect some sort of structure
# for the variables - or perhaps a single variable.
self.generic_visit(source)
v_rep_not_norm = self.as_sequence(source)
# What we have is a sequence of the data values we want to fill. The iterator at play
# here is the scope we want to use to run our Fill() calls to the TTree.
scope_fill = v_rep_not_norm.iterator_value().scope()
# Clean the data up so it is uniform and the next bit can proceed smoothly.
# If we don't have a tuple of data to log, turn it into a tuple.
seq_values = v_rep_not_norm.sequence_value()
if not isinstance(seq_values, crep.cpp_tuple):
seq_values = crep.cpp_tuple((v_rep_not_norm.sequence_value(),), scope_fill)
# Make sure the number of items is the same as the number of columns specified.
if len(seq_values.values()) != len(column_names):
raise Exception("Number of columns ({0}) is not the same as labels ({1}) in TTree creation".format(len(seq_values.values()), len(column_names)))
# Next, look at each on in turn to decide if it is a vector or a simple variable.
# Create a variable that we will fill for each one.
var_names = [(name, crep.cpp_variable(unique_name(name, is_class_var=True), self._gc.current_scope(), cpp_type=get_ttree_type(rep)))
for name, rep in zip(column_names, seq_values.values())]
# For each incoming variable, we need to declare something we are going to write.
for cv in var_names:
self._gc.declare_class_variable(cv[1])
# Next, emit the booking code
self._gc.add_book_statement(statement.book_ttree(tree_name, var_names))
# Note that the output file and tree are what we are going to return.
# The output filename is fixed - the hose code in AnalysisBase has that hard coded.
# To allow it to be different we have to modify that template too, and pass the
# information there. If more than one tree is written, the current code would
# lead to a bug.
node.rep = rh.cpp_ttree_rep("ANALYSIS.root", tree_name, self._gc.current_scope()) # type: ignore
# For each varable we need to save, cache it or push it back, depending.
# Make sure that it happens at the proper scope, where what we are after is defined!
s_orig = self._gc.current_scope()
for e_rep, e_name in zip(seq_values.values(), var_names):
scope_fill = self.code_fill_ttree(e_rep, e_name[1], scope_fill)
# The fill statement. This should happen at the scope where the tuple was defined.
# The scope where this should be done is a bit tricky (note the update above):
# - If a sequence, you want it where the sequence iterator is defined - or outside that scope
# - If a value, you want it at the level where the value is set.
self._gc.set_scope(scope_fill)
self._gc.add_statement(statement.ttree_fill(tree_name))
for e in zip(seq_values.values(), var_names):
if rep_is_collection(e[0]):
self._gc.add_statement(statement.container_clear(e[1][1]))
# And we are a terminal, so pop off the block.
self._gc.set_scope(s_orig)
self._gc.pop_scope()
return node.rep # type: ignore