def log(name, node):
return Assign(targets=[
Subscript(value=Subscript(value=Name(id=Log.log_dict, ctx=Load()),
slice=Index(value=Log.section_name),
ctx=Load()),
slice=Index(value=Str(s=name)),
ctx=Store())
],
value=node)
def compile_factory(fff: FlatFunctionFactory):
arguments = [*fff.arguments.keys()]
funname = fff.funname
unpacking = []
for i, (arg_group_name, arg_group) in enumerate(fff.arguments.items()):
for pos, sym in enumerate(arg_group):
rhs = Subscript(value=Name(id=arg_group_name, ctx=Load()),
slice=Index(Num(pos)),
ctx=Load())
val = Assign(targets=[Name(id=sym, ctx=Store())], value=rhs)
unpacking.append(val)
body = []
for (k, neq) in fff.preamble.items():
val = parse_string(neq).value
line = Assign(targets=[Name(id=k, ctx=Store())], value=val)
body.append(line)
for n, (k, neq) in enumerate(fff.content.items()):
# should the result of parse_string always of type Expr ?
val = parse_string(neq).value
line = Assign(targets=[Name(id=k, ctx=Store())], value=val)
body.append(line)
#
for n, (lhs, neq) in enumerate(fff.content.items()):
line = Assign(targets=[
Subscript(value=Name(id='out', ctx=Load()),
slice=Index(Num(n)),
ctx=Store())
],
value=Name(id=lhs, ctx=Load()))
body.append(line)
f = FunctionDef(name=funname,
args=ast_arguments(args=[arg(arg=a) for a in arguments] +
[arg(arg='out')],
vararg=None,
kwarg=None,
kwonlyargs=[],
kw_defaults=[],
defaults=[]),
body=unpacking + body,
decorator_list=[])
mod = Module(body=[f])
mmod = ast.fix_missing_locations(mod)
return mmod
def visit_Lambda(self, node):
""" Instrument a lambda expression by displaying the parameter values.
We create calls to trace assignment to each argument, then wrap them
all in a tuple together with the original expression, and pull the
original expression out of the tuple.
"""
new_node = self.generic_visit(node)
line_numbers = set()
self._find_line_numbers(new_node, line_numbers)
# trace lambda argument values
calls = [getattr(self._trace_assignment(target, node.lineno),
'value',
None)
for target in new_node.args.args
if getattr(target, 'id', 'self') != 'self' or
getattr(target, 'arg', 'self') != 'self']
args = [Num(n=min(line_numbers)),
Num(n=max(line_numbers))]
calls.insert(0, self._create_context_call('start_block', args).value)
calls.append(new_node.body)
new_node.body = Subscript(value=Tuple(elts=calls,
ctx=Load()),
slice=Index(value=Num(n=-1)),
ctx=Load())
return new_node
def compile_translationnode(self, srcnode, parent):
translated = Call(func=LoadName('_'),
args=[Str(srcnode.get_msgstr())],
starargs=None,
kwargs=None,
keywords=[])
named_children = [(name, node)
for name, node in srcnode.named_children()
if name is not None]
if not named_children:
# Simple case - no dynamic children for placeholder replacement
parent.body.append(Expr(value=Yield(translated)))
return
parent.body.append(
Assign(targets=[StoreName('__piglet_places')],
value=Dict([], []))
)
for name, node in named_children:
with self.collect_output(parent) as ACC:
self._compile(node, parent)
parent.body.append(
Assign(targets=[Subscript(value=LoadName('__piglet_places'),
slice=Index(value=Str(name)),
ctx=Store())],
value=Call(func=Attribute(value=Str(s=''), attr='join', ctx=Load()),
args=[LoadName(ACC)],
starargs=None,
kwargs=None,
keywords=[]))
)
for name, node in named_children:
translated = Call(
func=Attribute(value=translated, attr='replace', ctx=Load()),
args=[Str('${{{}}}'.format(name)),
Subscript(value=LoadName('__piglet_places'),
slice=Index(value=Str(name)),
ctx=Load())],
starargs=None,
kwargs=None,
keywords=[])
set_pos(translated, srcnode)
parent.body.append(Expr(value=Yield(translated)))
def visit_FunctionDef(self, node):
""" Instrument a function definition by creating a new report builder
for this stack frame and putting it in a local variable. The local
variable has the same name as the global variable so all calls can
use the same CONTEXT_NAME symbol, but it means that I had to use this:
x = globals()['x'].start_frame()
Kind of ugly, but I think it was worth it to handle recursive calls.
"""
if node.name == '__repr__':
return node
new_node = self.generic_visit(node)
line_numbers = set()
self._find_line_numbers(new_node, line_numbers)
first_line_number = min(line_numbers)
last_line_number = max(line_numbers)
args = [Num(n=first_line_number),
Num(n=last_line_number)]
try_body = new_node.body
globals_call = Call(func=Name(id='globals', ctx=Load()),
args=[],
keywords=[],
starargs=None,
kwargs=None)
global_context = Subscript(value=globals_call,
slice=Index(value=Str(s=CONTEXT_NAME)),
ctx=Load())
start_frame_call = Call(func=Attribute(value=global_context,
attr='start_frame',
ctx=Load()),
args=args,
keywords=[],
starargs=None,
kwargs=None)
context_assign = Assign(targets=[Name(id=CONTEXT_NAME, ctx=Store())],
value=start_frame_call)
new_node.body = [context_assign]
# trace function parameter values
for target in new_node.args.args:
if isinstance(target, Name) and target.id == 'self':
continue
if arg and isinstance(target, arg) and target.arg == 'self':
continue
new_node.body.append(self._trace_assignment(target, node.lineno))
handler_body = [self._create_context_call('exception'),
Raise()]
new_node.body.append(
TryExcept(body=try_body,
handlers=[ExceptHandler(body=handler_body)],
orelse=[],
finalbody=[]))
self._set_statement_line_numbers(try_body, first_line_number)
self._set_statement_line_numbers(handler_body, last_line_number)
return new_node
def construct_assign_dict(target, dict_name):
# construct_assign_dict('a','d')
# -> a = d['a']
from ast import Assign, Name, Store, Subscript, Load, Index, Str
a = target
d = dict_name
expr = Assign(targets=[Name(id=a, ctx=Store())],
value=Subscript(value=Name(id=d, ctx=Load()),
slice=Index(value=Str(s=a)),
ctx=Load()))
return expr
def set_slice(node):
"""
In Python 3.9 there's a new ast parser (PEG) that no longer wraps things in Index.
This function handles this issue.
:param node: An AST node
:type node: ```ast.AST```
:return: Original node, possibly wrapped in an ```Index```
:rtype: ```Union[ast.AST, Index]```
"""
return node if PY_GTE_3_9 else Index(node)
def construct_assign_array(target, array_name, index):
# construct_assign_array('a', 'v', i)
# -> a = v[i]
from ast import Assign, Name, Store, Subscript, Load, Index, Str, Num
a = target
d = array_name
i = index
expr = Assign(targets=[Name(id=a, ctx=Store())],
value=Subscript(value=Name(id=array_name, ctx=Load()),
slice=Index(value=Num(n=i)),
ctx=Load()))
return expr
def _generate_assignment((arg_name, arg_resource_handle)):
# type: (Tuple[basestring, int]) -> If
"""
We have a function that looks like:
def do_something(param, model_=INJECTED):
<...>
We insert into its beginning a statement like
___INJECT_CONTEXT_INTERNAL_RESOURCES = ___INJECT_CONTEXT_INTERNAL.resources
if model_ is INJECTED:
model_ = ___INJECT_CONTEXT_INTERNAL_RESOURCES[3]
if model is ___INJECT_CONTEXT_INTERNAL: # means that no resource is available
___INJECT_CONTEXT_INTERNAL_RESOURCES.flag_missing('model_')
Code outside of this function sets a global variable _INJECTED__model to point at the right thing.
"""
target_attribute = Subscript(
value=Name(id=INTERNAL_RESOURCES_NAME, ctx=Load()),
slice=Index(value=Num(n=arg_resource_handle)),
ctx=Load())
consequence = [
Assign(targets=[Name(id=arg_name, ctx=Store())],
value=target_attribute),
If(test=Compare(
left=Name(id=arg_name, ctx=Load()),
ops=[Is()],
comparators=[Name(id=INTERNAL_CONTEXT_NAME, ctx=Load())]),
body=[
Expr(value=Call(func=Attribute(value=Name(
id=INTERNAL_CONTEXT_NAME, ctx=Load()),
attr='flag_missing',
ctx=Load()),
keywords=[],
starargs=None,
kwargs=None,
args=[Str(s=arg_name)]))
],
orelse=[])
] # type: List[Union[Assign, If]
return If(test=Compare(
left=Name(id=arg_name, ctx=Load()),
ops=[Is()],
comparators=[default_nodes_mapping[arg_name]]),
body=consequence,
orelse=[])
def subscriptlist(self, s):
s = s[0]
subscript = None
if isinstance(s, Tree):
subscript = s.children
if len(subscript) == 1:
subscript = Index(value=subscript[0])
elif len(subscript) == 2:
subscript = Slice(lower=subscript[0],
upper=subscript[1],
step=None)
elif len(subscript) == 3:
subscript = Slice(lower=subscript[0],
upper=subscript[1],
step=subscript[2])
return Tree(data='subscriptlist', children=[subscript])
def it2literal(it):
"""
Convert a collection of constants into a type annotation
:param it: collection of constants
:type it: ```Union[Tuple[Union[str, int, float], ...], List[Union[str, int, float], ...]]```
:return: Subscript Literal for annotation
:rtype: ```Subscript```
"""
return Subscript(
Name("Literal", Load()),
Index(
value=Tuple(ctx=Load(), elts=list(map(set_value, it)), expr=None)
if len(it) > 1
else set_value(it[0])
),
Load(),
)
def object_instance_from_res(object_name: str, object_id: str,
cls_name: str) -> AnnAssign:
try:
is_valid_identifier(object_name)
except Arcor2Exception as e:
raise Arcor2Exception(f"Object name {object_name} invalid. {str(e)}")
try:
is_valid_type(cls_name)
except Arcor2Exception as e:
raise Arcor2Exception(f"Class name {cls_name} invalid. {str(e)}")
return AnnAssign(
target=Name(id=object_name, ctx=Store()),
annotation=Name(id=cls_name, ctx=Load()),
value=Subscript(value=get_name_attr("res", "objects"),
slice=Index(value=Str(s=object_id, kind="")),
ctx=Load()),
simple=1,
)
def compile_factory(fff: FlatFunctionFactory):
arguments = [*fff.arguments.keys()]
funname = fff.funname
unpacking = []
for i, (arg_group_name, arg_group) in enumerate(fff.arguments.items()):
for pos, sym in enumerate(arg_group):
rhs = Subscript(
value=Name(id=arg_group_name, ctx=Load()),
slice=Index(Num(pos)),
ctx=Load(),
)
val = Assign(targets=[Name(id=sym, ctx=Store())], value=rhs)
unpacking.append(val)
body = []
for (k, neq) in fff.preamble.items():
tree = parse_string(neq)
val = tree_to_ast(tree).value
line = Assign(targets=[Name(id=k, ctx=Store())], value=val)
body.append(line)
for n, (k, neq) in enumerate(fff.content.items()):
tree = parse_string(neq)
val = tree_to_ast(tree).value
line = Assign(targets=[Name(id=k, ctx=Store())], value=val)
body.append(line)
#
for n, (lhs, neq) in enumerate(fff.content.items()):
line = Assign(
targets=[
Subscript(value=Name(id="out", ctx=Load()),
slice=Index(Num(n)),
ctx=Store())
],
value=Name(id=lhs, ctx=Load()),
)
body.append(line)
if sys.version_info >= (3, 8, 0):
f = FunctionDef(
name=funname,
args=ast_arguments(
posonlyargs=[],
args=[arg(arg=a) for a in arguments] + [arg(arg="out")],
vararg=None,
kwarg=None,
kwonlyargs=[],
kw_defaults=[],
defaults=[],
),
body=unpacking + body,
decorator_list=[],
)
mod = Module(body=[f], type_ignores=[])
else:
f = FunctionDef(
name=funname,
args=ast_arguments(
args=[arg(arg=a) for a in arguments] + [arg(arg="out")],
vararg=None,
kwarg=None,
kwonlyargs=[],
kw_defaults=[],
defaults=[],
),
body=unpacking + body,
decorator_list=[],
)
mod = Module(body=[f])
mmod = ast.fix_missing_locations(mod)
return mmod
def visit_FunctionDef(self, node):
""" Instrument a function definition by creating a new report builder
for this stack frame and putting it in a local variable. The local
variable has the same name as the global variable so all calls can
use the same CONTEXT_NAME symbol, but it means that I had to use this:
x = globals()['x'].start_frame()
Kind of ugly, but I think it was worth it to handle recursive calls.
"""
new_node = self.generic_visit(node)
line_numbers = set()
find_line_numbers(new_node, line_numbers)
first_line_number = min(line_numbers)
last_line_number = max(line_numbers)
args = [Num(n=first_line_number), Num(n=last_line_number)]
start_frame_keywords = []
for decorator in new_node.decorator_list:
if getattr(decorator, 'id', None) == 'traced':
start_frame_keywords.append(
keyword(arg='is_decorated',
value=Name(id='True', ctx=Load())))
try_body = new_node.body
globals_call = Call(func=Name(id='globals', ctx=Load()),
args=[],
keywords=[],
starargs=None,
kwargs=None)
global_context = Subscript(value=globals_call,
slice=Index(value=Str(s=CONTEXT_NAME)),
ctx=Load())
start_frame_call = Call(func=Attribute(value=global_context,
attr='start_frame',
ctx=Load()),
args=args,
keywords=start_frame_keywords,
starargs=None,
kwargs=None)
context_assign = Assign(targets=[Name(id=CONTEXT_NAME, ctx=Store())],
value=start_frame_call)
new_node.body = [context_assign]
if isinstance(try_body[0], Expr) and isinstance(
try_body[0].value, Str):
# Move docstring back to top of function.
# noinspection PyUnresolvedReferences
new_node.body.insert(0, try_body.pop(0))
# trace function parameter values
arg_nodes = []
arg_nodes.extend(getattr(new_node.args, 'posonlyargs', []))
arg_nodes.extend(new_node.args.args)
arg_nodes.append(new_node.args.kwarg)
arg_nodes.append(new_node.args.vararg)
arg_nodes.extend(new_node.args.kwonlyargs)
for target in arg_nodes:
if target is None:
continue
if isinstance(target, Name) and target.id == 'self':
continue
if isinstance(target, arg) and target.arg == 'self':
continue
new_node.body.append(self._trace_assignment(target, node.lineno))
if try_body:
handler_body = [self._create_context_call('exception'), Raise()]
new_node.body.append(
Try(body=try_body,
handlers=[ExceptHandler(body=handler_body)],
orelse=[],
finalbody=[]))
self._set_statement_line_numbers(try_body, first_line_number)
self._set_statement_line_numbers(handler_body, last_line_number)
return new_node
def visit_Name(self, node, value):
return copy_location(
Subscript(value=cte_to_node(value),
slice=Index(value=Str(s=node.id)),
ctx=node.ctx), node)
def __assign_var(self, lineno, var, loop=False): lineno_key = Subscript(value=Name(id=Environment.name,ctx=Load()), slice=Index(value=Num(n=lineno)), ctx=Load()) var_key = lineno_key if not loop else Subscript(value=lineno_key, slice=Index(value=Name(id=Environment.iter_num, ctx=Load())), ctx=Load()) return Assign(targets=[Subscript(value=var_key, slice=Index(value=Str(s=var)), ctx=Store())], value=Name(id=var,ctx=Load()))
def __delete_var(self, lineno, var): return Delete(targets=[Subscript(value=Subscript(value=Name(id=Environment.name,ctx=Load()), slice=Index(value=Num(n=lineno)), ctx=Load()), slice=Index(value=Str(s=var)), ctx=Del())])
def ast_index(name: str):
return Index(value=Name(id=name, ctx=Load()))
def init(self, lineno, args): init_call = Assign(targets=[Subscript(value=Name(id=Environment.name,ctx=Load()), slice=Index(value=Num(n=lineno)), ctx=Store())], value=Dict()) param_init = self.assign(lineno,args) self.lineno = lineno+1 return [init_call] + param_init
annotation=Name(
"str",
Load(),
),
arg="tfds_dir",
),
set_arg(
annotation=Subscript(
Name(
"Literal",
Load(),
),
Index(value=Tuple(
ctx=Load(),
elts=[
set_value("np", ),
set_value("tf", ),
],
expr=None,
)),
Load(),
),
arg="K",
),
set_arg(
annotation=Subscript(
Name(
"Optional",
Load(),
),
Index(value=Name("bool", Load())),
Load(),
def index(x):
Index(Num(x))
import ast
import os
from ast import AnnAssign, Assign, ClassDef, Constant, ImportFrom, Index, List, Load, Module, Name, Store, Subscript, \
alias
from ast_decompiler import decompile
dataclass_decorator = Name(id='dataclass', ctx=Load())
text_data_value = Subscript(
value=Name(id='Optional', ctx=Load()),
slice=Index(value=Name(id='TextData', ctx=Load())),
ctx=Load()
)
chart_xy_data_value = Subscript(
value=Name(id='Optional', ctx=Load()),
slice=Index(value=Name(id='ChartXYData', ctx=Load())),
ctx=Load()
)
chart_category_data_value = Subscript(
value=Name(id='Optional', ctx=Load()),
slice=Index(value=Name(id='ChartCategoryData', ctx=Load())),
ctx=Load()
)
chart_bubble_data_value = Subscript(
value=Name(id='Optional', ctx=Load()),
slice=Index(value=Name(id='ChartBubbleData', ctx=Load())),
ctx=Load()
)
table_data_value = Subscript(
value=Name(id='Optional', ctx=Load()),
slice=Index(value=Name(id='TableData', ctx=Load())),
def compile_function_ast(expressions, symbols, arg_names, output_names=None, funname='anonymous', return_ast=False, print_code=False, definitions=None, vectorize=True, use_file=False):
'''
expressions: list of equations as string
'''
from collections import OrderedDict
table = OrderedDict()
aa = arg_names
if output_names is not None:
aa = arg_names + [output_names]
for a in aa:
symbol_group = a[0]
date = a[1]
an = a[2]
for b in symbols[symbol_group]:
index = symbols[symbol_group].index(b)
table[(b, date)] = (an, index)
table_symbols = {k: (std_date_symbol(*k)) for k in table.keys()}
# standard assignment: i.e. k = s[0]
index = lambda x: Index(Num(x))
# declare symbols
aux_short_names = [e[2] for e in arg_names if e[0]=='auxiliaries']
preamble = []
for k in table: # order it
# k : var, date
arg, pos = table[k]
if not (arg in aux_short_names):
std_name = table_symbols[k]
val = Subscript(value=Name(id=arg, ctx=Load()), slice=index(pos), ctx=Load())
line = Assign(targets=[Name(id=std_name, ctx=Store())], value=val)
if arg != 'out':
preamble.append(line)
body = []
std_dates = StandardizeDates(symbols, aa)
if definitions is not None:
for k,v in definitions.items():
if isinstance(k, str):
lhs = ast.parse(k).body[0].value
if isinstance(v, str):
rhs = ast.parse(v).body[0].value
else:
rhs = v
lhs = std_dates.visit(lhs)
rhs = std_dates.visit(rhs)
vname = lhs.id
line = Assign(targets=[Name(id=vname, ctx=Store())], value=rhs)
preamble.append(line)
outs = []
for i, expr in enumerate(expressions):
expr = ast.parse(expr).body[0].value
# if definitions is not None:
# expr = ReplaceName(defs).visit(expr)
rexpr = std_dates.visit(expr)
rhs = rexpr
if output_names is not None:
varname = symbols[output_names[0]][i]
date = output_names[1]
out_name = table_symbols[(varname, date)]
else:
out_name = 'out_{}'.format(i)
line = Assign(targets=[Name(id=out_name, ctx=Store())], value=rhs)
body.append(line)
line = Assign(targets=[Subscript(value=Name(id='out', ctx=Load()),
slice=index(i), ctx=Store())], value=Name(id=out_name, ctx=Load()))
body.append(line)
arg_names = [e for e in arg_names if e[0]!="auxiliaries"]
args = [e[2] for e in arg_names] + ['out']
if is_python_3:
from ast import arg
f = FunctionDef(name=funname, args=arguments(args=[arg(arg=a) for a in args], vararg=None, kwarg=None, kwonlyargs=[], kw_defaults=[], defaults=[]),
body=preamble + body, decorator_list=[])
else:
f = FunctionDef(name=funname, args=arguments(args=[Name(id=a, ctx=Param()) for a in args], vararg=None, kwarg=None, kwonlyargs=[], kw_defaults=[], defaults=[]),
body=preamble + body, decorator_list=[])
mod = Module(body=[f])
mod = ast.fix_missing_locations(mod)
if print_code:
s = "Function {}".format(mod.body[0].name)
print("-" * len(s))
print(s)
print("-" * len(s))
print(to_source(mod))
if vectorize:
from numba import float64, void
coredims = [len(symbols[an[0]]) for an in arg_names]
signature = str.join(',', ['(n_{})'.format(d) for d in coredims])
n_out = len(expressions)
if n_out in coredims:
signature += '->(n_{})'.format(n_out)
# ftylist = float64[:](*([float64[:]] * len(coredims)))
fty = "void(*[float64[:]]*{})".format(len(coredims)+1)
else:
signature += ',(n_{})'.format(n_out)
fty = "void(*[float64[:]]*{})".format(len(coredims)+1)
ftylist = [fty]
else:
signature=None
ftylist=None
if use_file:
fun = eval_ast_with_file(mod, print_code=True)
else:
fun = eval_ast(mod)
jitted = njit(fun)
if vectorize:
gufun = guvectorize([fty], signature, target='parallel', nopython=True)(fun)
return jitted, gufun
else:
return jitted
def make_function(equations,
arguments,
parameters,
targets=None,
rhs_only=False,
definitions={},
funname='anonymous'):
compat = lambda s: s.replace("^", "**").replace('==', '=').replace(
'=', '==')
equations = [compat(eq) for eq in equations]
if isinstance(arguments, list):
arguments = OrderedDict([('arg_{}'.format(i), k)
for i, k in enumerate(arguments)])
## replace = by ==
known_variables = [a[0] for a in sum(arguments.values(), [])]
known_definitions = [a for a in definitions.keys()]
known_parameters = [a[0] for a in parameters]
all_variables = known_variables + known_definitions
known_functions = []
known_constants = []
if targets is not None:
all_variables.extend([o[0] for o in targets])
targets = [std_tsymbol(o) for o in targets]
else:
targets = ['_out_{}'.format(n) for n in range(len(equations))]
all_symbols = all_variables + known_parameters
equations = [parse(eq) for eq in equations]
definitions = {k: parse(v) for k, v in definitions.items()}
defs_incidence = {}
for sym, val in definitions.items():
cn = CountNames(known_definitions, [], [])
cn.visit(val)
defs_incidence[(sym, 0)] = cn.variables
# return defs_incidence
from dolo.compiler.codegen import to_source
equations_incidence = {}
to_be_defined = set([])
for i, eq in enumerate(equations):
cn = CountNames(all_variables, known_functions, known_constants)
cn.visit(eq)
equations_incidence[i] = cn.variables
to_be_defined = to_be_defined.union(
[a for a in cn.variables if a[0] in known_definitions])
deps = []
for tv in to_be_defined:
ndeps = get_deps(defs_incidence, tv)
deps.extend(ndeps)
deps = [d for d in unique(deps)]
sds = StandardizeDatesSimple(all_symbols)
new_definitions = OrderedDict()
for k in deps:
val = definitions[k[0]]
nval = timeshift(val, all_variables, k[1]) # function to print
# dprint(val)
new_definitions[std_tsymbol(k)] = sds.visit(nval)
new_equations = []
for n, eq in enumerate(equations):
d = match(parse("_x == _y"), eq)
if d is not False:
lhs = d['_x']
rhs = d['_y']
if rhs_only:
val = rhs
else:
val = ast.BinOp(left=rhs, op=Sub(), right=lhs)
else:
val = eq
new_equations.append(sds.visit(val))
# preambleIndex(Num(x))
preamble = []
for i, (arg_group_name, arg_group) in enumerate(arguments.items()):
for pos, t in enumerate(arg_group):
sym = std_tsymbol(t)
rhs = Subscript(value=Name(id=arg_group_name, ctx=Load()),
slice=Index(Num(pos)),
ctx=Load())
val = Assign(targets=[Name(id=sym, ctx=Store())], value=rhs)
preamble.append(val)
for pos, p in enumerate(parameters):
sym = std_tsymbol(p)
rhs = Subscript(value=Name(id='p', ctx=Load()),
slice=Index(Num(pos)),
ctx=Load())
val = Assign(targets=[Name(id=sym, ctx=Store())], value=rhs)
preamble.append(val)
# now construct the function per se
body = []
for k, v in new_definitions.items():
line = Assign(targets=[Name(id=k, ctx=Store())], value=v)
body.append(line)
for n, neq in enumerate(new_equations):
line = Assign(targets=[Name(id=targets[n], ctx=Store())],
value=new_equations[n])
body.append(line)
for n, neq in enumerate(new_equations):
line = Assign(targets=[
Subscript(value=Name(id='out', ctx=Load()),
slice=Index(Num(n)),
ctx=Store())
],
value=Name(id=targets[n], ctx=Load()))
body.append(line)
from ast import arg, FunctionDef, Module
from ast import arguments as ast_arguments
from dolo.compiler.function_compiler_ast import to_source
f = FunctionDef(
name=funname,
args=ast_arguments(args=[arg(arg=a) for a in arguments.keys()] +
[arg(arg='p'), arg(arg='out')],
vararg=None,
kwarg=None,
kwonlyargs=[],
kw_defaults=[],
defaults=[]),
body=preamble + body,
decorator_list=[])
mod = Module(body=[f])
mod = ast.fix_missing_locations(mod)
return mod
expr_target=None,
),
AnnAssign(
annotation=Subscript(
Name(
"Literal",
Load(),
),
Index(
value=Tuple(
elts=list(
map(
set_value,
(
"np",
"tf",
),
)
),
ctx=Load(),
expr=None,
)
),
Load(),
),
simple=1,
target=Name("K", Store()),
value=set_value("np"),
expr=None,
expr_target=None,
expr_annotation=None,
def _hoist_variables_to_piglet_context(astnode, exclude_functions=frozenset()):
"""
Template functions extract all local variables from the
:var:`piglet.runtime.data` thread local
"""
# Names we never hoist.
# "None" would raise a SyntaxError if we try assigning to it. The others
# are used internally by piglet and need to be reserved.
restricted_names = {'None', 'True', 'False',
'value_of', 'defined', 'Markup', 'iter',
'AttributeError', 'Exception',
'print', 'getattr'}
# Mapping of function -> names used within function
func_names = {}
# Mapping of {function name: {<ancestor nodes>, ...}}
func_ancestors = {}
# All names discovered together with their ast location
names = []
for fn, ancestors in _get_function_defs(astnode):
func_ancestors.setdefault(fn.name, set()).add(ancestors)
func_names[fn] = set()
is_reserved = lambda n: n.id in restricted_names
is_piglet = lambda n: n.id.startswith('__piglet')
is_function = lambda n, a: any(
a[:len(fancestors)] == fancestors
for fancestors in func_ancestors.get(n.id, set()))
names = ((ancestors, node)
for node, ancestors in _get_out_of_scope_names(astnode)
if not (is_reserved(node) or
is_piglet(node) or
is_function(node, ancestors)))
for ancestors, node in names:
container_func = next((a
for a in reversed(ancestors)
if isinstance(a, FunctionDef)), None)
if container_func in exclude_functions:
continue
if container_func is None:
lineno = getattr(node, 'lineno', '(unknown line number)')
raise AssertionError("Unexpected variable found at {}: {}"
.format(lineno, node.id))
func_names[container_func].add(node.id)
for f, names in func_names.items():
assignments = [
Assign(targets=[StoreName('__piglet_ctx')],
value=Subscript(
value=LoadAttribute('__piglet_rtdata.context'),
slice=Index(value=Num(n=-1)),
ctx=Load()
))
]
for n in sorted(names):
is_builtin = n in builtin_names
if is_builtin:
default = LoadAttribute(
LoadAttribute('__piglet_rt', 'builtins'), n)
else:
default = Call(func=Attribute(value=Name(id='__piglet_rt',
ctx=Load()),
attr='Undefined',
ctx=Load()),
args=[Str(s=n)],
starargs=None,
kwargs=None,
keywords=[])
value = Call(func=Attribute(value=Name(id='__piglet_ctx',
ctx=Load()),
attr='get',
ctx=Load()),
args=[Str(s=n), default],
starargs=None,
kwargs=None,
keywords=[])
a = Assign(targets=[Name(id=n, ctx=Store())], value=value)
assignments.append(a)
f.body[0:0] = assignments
return astnode
annotation=Name(
"str",
Load(),
),
arg="tfds_dir",
),
set_arg(
annotation=Subscript(
Name(
"Literal",
Load(),
),
Index(value=Tuple(
ctx=Load(),
elts=[
set_value("np", ),
set_value("tf", ),
],
expr=None,
)),
Load(),
),
arg="K",
),
set_arg(
annotation=Subscript(
Name(
"Optional",
Load(),
),
Index(value=Name("bool", Load())),
Load(),
def compile_function_ast(expressions,
symbols,
arg_names,
output_names=None,
funname='anonymous',
data_order='columns',
use_numexpr=False,
return_ast=False,
print_code=False,
definitions=None):
'''
expressions: list of equations as string
'''
vectorization_type = 'ellipsis'
data_order = 'columns'
from collections import OrderedDict
table = OrderedDict()
aa = arg_names
if output_names is not None:
aa = arg_names + [output_names]
for a in aa:
symbol_group = a[0]
date = a[1]
an = a[2]
for b in symbols[symbol_group]:
index = symbols[symbol_group].index(b)
table[(b, date)] = (an, index)
table_symbols = {k: (std_date_symbol(*k)) for k in table.keys()}
if data_order is None:
# standard assignment: i.e. k = s[0]
index = lambda x: Index(Num(x))
elif vectorization_type == 'ellipsis':
el = Ellipsis()
if data_order == 'columns':
# column broadcasting: i.e. k = s[...,0]
index = lambda x: ExtSlice(dims=[el, Index(value=Num(n=x))])
else:
# rows broadcasting: i.e. k = s[0,...]
index = lambda x: ExtSlice(dims=[Index(value=Num(n=x)), el])
# declare symbols
preamble = []
for k in table: # order it
# k : var, date
arg, pos = table[k]
std_name = table_symbols[k]
val = Subscript(value=Name(id=arg, ctx=Load()),
slice=index(pos),
ctx=Load())
line = Assign(targets=[Name(id=std_name, ctx=Store())], value=val)
preamble.append(line)
if use_numexpr:
for i in range(len(expressions)):
# k : var, date
val = Subscript(value=Name(id='out', ctx=Load()),
slice=index(i),
ctx=Load())
line = Assign(targets=[Name(id='out_{}'.format(i), ctx=Store())],
value=val)
preamble.append(line)
body = []
std_dates = StandardizeDates(symbols, aa)
if definitions is not None:
defs = {
e: ast.parse(definitions[e]).body[0].value
for e in definitions
}
for i, expr in enumerate(expressions):
expr = ast.parse(expr).body[0].value
if definitions is not None:
expr = ReplaceName(defs).visit(expr)
rexpr = std_dates.visit(expr)
if not use_numexpr:
rhs = rexpr
else:
import codegen
src = codegen.to_source(rexpr)
rhs = Call(func=Name(id='evaluate', ctx=Load()),
args=[Str(s=src)],
keywords=[
keyword(arg='out',
value=Name(id='out_{}'.format(i),
ctx=Load()))
],
starargs=None,
kwargs=None)
if not use_numexpr:
val = Subscript(value=Name(id='out', ctx=Load()),
slice=index(i),
ctx=Store())
line = Assign(targets=[val], value=rhs)
else:
line = Expr(
value=rhs
) #Assign(targets=[Name(id='out_{}'.format(i), ctx=Load())], value=rhs )
body.append(line)
if output_names is not None:
varname = symbols[output_names[0]][i]
date = output_names[1]
out_name = table_symbols[(varname, date)]
line = Assign(targets=[Name(id=out_name.format(i), ctx=Store())],
value=Name(id='out_{}'.format(i), ctx=Store()))
# body.append(line)
args = [e[2] for e in arg_names] + ['out']
f = FunctionDef(name=funname,
args=arguments(
args=[Name(id=a, ctx=Param()) for a in args],
vararg=None,
kwarg=None,
defaults=[]),
body=preamble + body,
decorator_list=[])
mod = Module(body=[f])
mod = ast.fix_missing_locations(mod)
# print_code=True
if print_code:
s = "Function {}".format(mod.body[0].name)
print("-" * len(s))
print(s)
print("-" * len(s))
import codegen
print(codegen.to_source(mod))
if return_ast:
return mod
else:
fun = eval_ast(mod)
return fun
