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