def test_ast_tree_edit_distance() -> None: node1 = ast.parse("a=1") node2 = ast.parse("a=1") assert ast_tree_edit_distance(node1, node2) == 0 node2 = ast.parse("b=1") assert ast_tree_edit_distance(node1, node2) == 1 node2 = ast.parse("b=2") assert ast_tree_edit_distance(node1, node2) == 1.5 node1 = Assign( targets=[Name(id="a", ctx=Store())], value=Constant(value=1, kind=None), type_comment=None, ) node2 = Assign( targets=[Name(id="a", ctx=Store())], value=Constant(value=1, kind=None), type_comment=None, ) assert ast_tree_edit_distance(node1, node2) == 0 node1 = ast.parse("") node2 = ast.parse("") assert ast_tree_edit_distance(node1, node2) == 0 node1 = ast.parse("a") node2 = ast.parse("a") assert ast_tree_edit_distance(node1, node2) == 0 node1 = Expr(Name("a", Load())) node2 = Expr(Name("a", Load())) assert ast_tree_edit_distance(node1, node2) == 0 node1 = Name("a", Load()) node2 = Name("a", Load()) assert ast_tree_edit_distance(node1, node2) == 0
def add_method_call( body: List, instance: str, method: str, args: List, kwargs: List, returns: List[str], index: Optional[int] = None ) -> None: """Adds method call to be body of a container. By default, it appends. When index is specified, it inserts. :param body: :param instance: :param method: :param args: :param kwargs: :param index: :param returns: :return: """ call = Call(func=get_name_attr(instance, method), args=args, keywords=kwargs) cont: Union[Expr, Assign, None] = None if not returns: cont = Expr(value=call) elif len(returns) == 1: # TODO AnnAssign?? cont = Assign(targets=[Name(id=returns[0], ctx=Store())], value=call) else: cont = Assign(targets=[Tuple(elts=[Name(id=ret, ctx=Store()) for ret in returns], ctx=Store())], value=call) if index is None: body.append(cont) else: body.insert(index, cont)
def collect_output(self, parent): """ Context manager that collects any yield expressions added to ``parent`` and turns them into calls to ``__piglet_acc_list.append``. The name of the accumulator object is returned by the function """ acc = self.unique_id('acc') append = self.unique_id('append') pos = len(parent.body) parent.body.append(Assign(targets=[StoreName(acc)], value=List(elts=[], ctx=Load()))) parent.body.append(Assign(targets=[StoreName(append)], value=Attribute(value=LoadName(acc), attr='append', ctx=Load()))) yield acc for n in parent.body[pos:]: for node, ancestors in astwalk(n): if isinstance(node, Expr) and isinstance(node.value, Yield): node.value = Call(func=LoadName(append), args=[node.value.value], starargs=None, kwargs=None, keywords=[])
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 ast_get_last_expression(code : str): """ Modify code so that if the last statement is an "Expr" or "Await" statement, we return that into "EXEC-LAST-EXPRESSION" """ from ast import ( fix_missing_locations, parse, Assign, Await, Constant, Expr, Name, Store ) tree = parse(code) targets = [Name("EXEC-LAST-EXPRESSION", ctx = Store())] if isinstance(tree.body[-1], (Expr, Await)): tree.body[-1] = Assign(targets, tree.body[-1].value) else: tree.body.append(Assign(targets, Constant(None, None))) fix_missing_locations(tree) return tree
def collect_var_assignments(pat, val): """This helper function ensures that the pattern is used to properly assign all subfields of the given AST for use in the clause body E.g., for PatternConstructor(A, PatternVar(v), PatternWildcard(), PatternConstructor(B, PatternVar(w))) we would want to have v = a.fields[0] w = a.fields[2].fields[0] """ if isinstance(pat, relay.PatternWildcard): return [] if isinstance(pat, relay.PatternVar): return [Assign([self.include_var(pat.var, assign=True)], val)] # constructor pattern: assign each field of the value # based on subpatterns assignments = [] for i in range(len(pat.patterns)): # we want the assignments for val.fields[i] field = ast.Subscript( ast.Attribute(val, "fields", Load()), ast.Index(Num(i)), Load() ) assignments += collect_var_assignments(pat.patterns[i], field) return assignments
def _add_coverage_dummy_node(tree, macronode, macroname): '''Force `macronode` to be reported as covered by coverage tools. The dummy node will be injected to `tree`. The `tree` must appear in a position where `ast.NodeTransformer.visit` may return a list of nodes. `macronode` is the macro invocation node to copy source location info from. `macroname` is included in the dummy node, to ease debugging. ''' # `macronode` itself might be macro-generated. In that case don't bother. if not hasattr(macronode, 'lineno') and not hasattr( macronode, 'col_offset'): return tree if tree is None: tree = [] elif isinstance(tree, AST): tree = [tree] # The dummy node must actually run to get coverage, so an `ast.Pass` won't do. # It must *do* something, or CPython optimizes it away, so an `ast.Expr` won't do. # We must set location info manually, because we run after `expand`. v = copy_location( Constant( value=f"source line {macronode.lineno} invoked macro {macroname}"), macronode) t = copy_location(Name(id="_mcpyrate_coverage", ctx=Store()), macronode) dummy = copy_location(Assign(targets=[t], value=v), macronode) tree.insert( 0, Done(dummy) ) # mark as Done so any expansions further out won't mess this up. return tree
def test_annotate_ancestry(self) -> None: """Tests that `annotate_ancestry` properly decorates""" node = Module( body=[ AnnAssign( annotation=Name( "str", Load(), ), simple=1, target=Name("dataset_name", Store()), value=set_value("~/tensorflow_datasets"), expr=None, expr_target=None, expr_annotation=None, ), Assign( annotation=None, simple=1, targets=[Name("epochs", Store())], value=set_value("333"), expr=None, expr_target=None, expr_annotation=None, **maybe_type_comment ), ], stmt=None, ) self.assertFalse(hasattr(node.body[0], "_location")) self.assertFalse(hasattr(node.body[1], "_location")) annotate_ancestry(node) self.assertEqual(node.body[0]._location, ["dataset_name"]) self.assertEqual(node.body[1]._location, ["epochs"])
class ImportBombast(RenameBombast): # import sys -> sys = __import__('sys', globals(), locals(), [], 0) one = Transformation(lambda n: Assign( targets=[Name(id=n.names[0].name, ctx=Store())], value=Call(func=Name(id='__import__', ctx=Load()), args=[ Str(s=n.names[0].name), Call(func=Name(id='globals', ctx=Load()), args=[], keywords=[], starargs=None, kwargs=None), Call(func=Name(id='locals', ctx=Load()), args=[], keywords=[], starargs=None, kwargs=None), List(elts=[], ctx=Load()), Num(n=0) ], keywords=[], starargs=None, kwargs=None))) def transform(self): num_imports = len(self.node.names) if num_imports == 1: return self.one.transform(self.node) else: return self.node
def convert_output(ret_type): """Use the function return type to produce auxiliary variables to store outputs. Returns ([assignments of output vars], [extra arguments to pass to op call], expression collecting output)""" if isinstance(ret_type, relay.TensorType): output_var_name = self.generate_var_name("_out") output_var = Name(output_var_name, Load()) shape = ast.Tuple( [Num(dim) for dim in ret_type.concrete_shape], Load()) # create a new NDArray of the right shape and dtype assign_output = Assign( [Name(output_var_name, Store())], self.create_call("nd.array", [ self.create_call("numpy.empty", [shape, Str(ret_type.dtype)]) ]), ) return ([assign_output], [output_var], output_var) assert isinstance(ret_type, relay.TupleType) assignments = [] extra_args = [] fields = [] for t in ret_type.fields: inner_assignments, inner_args, inner_output = convert_output(t) assignments += inner_assignments extra_args += inner_args fields.append(inner_output) fields = [ast.List(fields, Load())] return (assignments, extra_args, self.create_call("_container.tuple_object", fields))
def s(self): """ id = expression | IF cond THEN s | IF cond THEN s else s | WHILE cond DO s """ token = self.current_token if token.type == IDN: varNode = Var(token) self.eat(IDN) op = self.current_token self.eat(EQL) right = self.expr() return Assign(varNode, op, right) elif token.type == IF: self.eat(IF) condition = self.cond() self.eat(THEN) statement1 = self.s() if self.current_token.type == ELSE: statement2 = self.s() return IfThenElse(condition, statement1, statement2) return IfThen(condition, statement1) elif token.type == WHILE: self.eat(WHILE) condition = self.cond() self.eat(DO) statement = self.s() return WhileDo(condition, statement)
def visit_match(self, match: Expr): """For matches, we wrap the entire expression in a thunk because it is easiest to implement them using if statements. For each clause, we generate a function that checks if the pattern matches. If yes, we call a function that assigns the variables appropriately and invokes the clause body.""" data, defs = self.visit(match.data) data_var = self.generate_var_name("_match_data") # must ensure the data clause is executed exactly once thunk_body = [Assign([Name(data_var, Store())], data)] for clause in match.clauses: check_expr = self.create_match_check(clause.lhs, Name(data_var, Load())) body_def, body_name = self.create_match_clause_body( clause.lhs, clause.rhs) defs.append(body_def) # equiv: if check(data): return body(data) thunk_body.append( ast.If(check_expr, [ Return( self.create_call(body_name, [Name(data_var, Load())])) ], [])) # finally if nothing matches we have a failed assert (should never happen) thunk_body.append( ast.Assert(NameConstant(False), Str("Match was not exhaustive"))) thunk_name = self.generate_function_name("_match_thunk") thunk_def = self.create_def(thunk_name, [], defs + thunk_body) return (self.create_call(thunk_name, []), [thunk_def])
def convert(self, prog: Expr): """This method converts the passed Relay expression into a Python AST object with equivalent semantics. The Python AST can be executed using exec(); it can be turned into text and inspected using astor. """ optimized = self.optimize(prog) # start with conversion prelude (imports) and convert global defs body = [] body += PROLOGUE body += self.convert_module() prog_body, extra_defs = self.visit(optimized) body += extra_defs # we finally must assign the final expression to the output var # so it can be read after running EXEC body.append(Assign([Name(OUTPUT_VAR_NAME, Store())], prog_body)) global __MAJOR__, __MINOR__ if __MAJOR__ == 3 and __MINOR__ == 8: return ast.fix_missing_locations( ast.Module(body=body, type_ignores=[])) else: return ast.fix_missing_locations(ast.Module(body=body))
def AugAssignToAssign(node): store_var = node.target load_var = Name(id=store_var.id, ctx=Load()) return Assign(targets=[store_var], value=BinOp(left=load_var, op=node.op, right=node.value))
def convert_output(ret_type): """Use the function return type to produce auxiliary variables to store outputs. Returns ([assignments of output vars], [extra arguments to pass to op call], expression collecting output)""" if isinstance(ret_type, relay.TensorType): output_var_name = self.generate_var_name('_out') output_var = Name(output_var_name, Load()) shape = ast.Tuple( [Num(dim) for dim in ret_type.concrete_shape], Load()) # create a new TensorValue of the right shape and dtype assign_output = Assign( [Name(output_var_name, Store())], self.create_call('TensorValue', [ self.create_call('numpy.empty', [shape, Str(ret_type.dtype)]) ])) # we pass the data field as an argument extra_arg = ast.Attribute(output_var, 'data', Load()) return ([assign_output], [extra_arg], output_var) assert isinstance(ret_type, relay.TupleType) assignments = [] extra_args = [] fields = [] for t in ret_type.fields: inner_assignments, inner_args, inner_output = convert_output(t) assignments += inner_assignments extra_args += inner_args fields.append(inner_output) return (assignments, extra_args, self.create_call('TupleValue', fields))
def assignment(self): #ASSIGNMENT : VAR ASSGN EXPR var_node = self.variable() token = self.current_token self.eat(Token_Type.ASSGN) expr = self.expr() return Assign(var_node,token,expr)
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 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 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 eval_if(self, test, line_state): test = Module(body=[ Assign(targets=[Name(id="__test_cond_result", ctx=Store())], value=test) ]) fix_missing_locations(test) code = compile(test, "<string>", "exec") exec code in line_state return line_state["__test_cond_result"]
def getGlobalVariable(name): # _lambda_0 return Assign( targets=[ Name( id=name + '_locals', ctx=Store() ) ], value=Dict(keys=[], values=[]), type_comment=None)
def make_slides_assign(slides_count): return Assign( targets=[Name(id='SLIDES', ctx=Store())], value=List( elts=[Name(id=f'Slide{idx}', ctx=Load()) for idx in range(1, slides_count + 1)], ctx=Load() ), type_comment=None )
def assign(p): ''' p[0] - LET p[1] - variable name p[2] - = p[3] - expressao ''' id = p[1].getstr() valor = expression(p) return Assign(self.builder, self.module, id, valor)
def test(): """ Test Function """ stmt = Program([ Function("f", [], [ For("i", 1, 10, 1, [ Assign("x", IntValue(3)), Assign("x", IntValue(4)), IfThenElse(BoolValue(True), []), While(BoolValue(True), [Return(IntValue(3))]), Assign("x", IntValue(5)) ]) ]), Function("g", [], [Assign("x", IntValue(3))]), Assign("x", IntValue(1)) ]) nodes, edges = gen_cfg_main(stmt) gen_dot(nodes, edges)
def add_builtins(self, fn): fn.body.extend([ Assign(targets=[Name(id='value_of', ctx=Store())], value=Attribute(value=Name(id='__piglet_ctx', ctx=Load()), attr='get', ctx=Load())), Assign(targets=[Name(id='defined', ctx=Store())], value=Attribute(value=Name(id='__piglet_ctx', ctx=Load()), attr='__contains__', ctx=Load())), Assign(targets=[Name(id='__piglet_escape', ctx=Store())], value=Attribute(value=Name(id='__piglet_rt', ctx=Load()), attr='escape', ctx=Load())), Assign(targets=[Name(id='__piglet_ustr', ctx=Store())], value=Attribute(value=Name(id='__piglet_rt', ctx=Load()), attr='ustr', ctx=Load())), ])
def document(tree, **kw): """ This macro takes literal strings and converts them into: _help_ID = type_hint+STRING where: ID is the first target of the last assignment. type_hint is the assigned type and default value (only works for a few types) STRING is the literal string """ for n in range(len(tree)): s = tree[n] if not n: prev = s continue # The whole sentence is a string? if (isinstance(s, Expr) and isinstance(s.value, Str) and # and the previous is an assign isinstance(prev, Assign)): # noqa: E128 # Apply it to the first target target = prev.targets[0] value = prev.value # Extract its name # variables and attributes are supported if isinstance(target, Name): name = target.id is_attr = False elif isinstance(target, Attribute): name = target.attr is_attr = True else: continue # Remove starting underscore if name[0] == '_': name = name[1:] # Create a _help_ID doc_id = '_help_' + name # Create the type hint for numbers, strings and booleans type_hint = '' if isinstance(value, Num): type_hint = '[number={}]'.format(value.n) elif isinstance(value, Str): type_hint = "[string='{}']".format(value.s) elif isinstance(value, NameConstant) and isinstance( value.value, bool): type_hint = '[boolean={}]'.format(str(value.value).lower()) # Transform the string into an assign for _help_ID if is_attr: target = Attribute(value=Name(id='self', ctx=Load()), attr=doc_id, ctx=Store()) else: target = Name(id=doc_id, ctx=Store()) tree[n] = Assign(targets=[target], value=Str(s=type_hint + s.value.s)) prev = s # Return the modified AST return tree
def make_assign(target, value): """ Return an assign node to assign value to target. """ return copy_loc( value, Assign( targets=[target], value=value, ), )
def document(tree, **kw): """ This macro takes literal strings and converts them into: _help_ID = type_hint+STRING where: ID is the first target of the last assignment. type_hint is the assigned type and default value (only works for a few types) STRING is the literal string """ # Simplify it just to show the problem isn't related to the content of the macro # Note: This triggers another issue, Expr nodes can be optimized out if not assigned to a target # return tree for n in range(len(tree)): s = tree[n] if not n: prev = s continue # The whole sentence is a string? if (isinstance(s, Expr) and isinstance(s.value, Str) and # and the previous is an assign isinstance(prev, Assign)): # noqa: E128 # Apply it to the first target target = prev.targets[0] value = prev.value # Extract its name # variables and attributes are supported if isinstance(target, Name): name = target.id is_attr = False elif isinstance(target, Attribute): name = target.attr is_attr = True # Create a _help_ID doc_id = '_help_'+name # Create the type hint for numbers, strings and booleans type_hint = '' if isinstance(value, Num): type_hint = '[number={}]'.format(value.n) elif isinstance(value, Str): type_hint = "[string='{}']".format(value.s) elif isinstance(value, NameConstant) and isinstance(value.value, bool): type_hint = '[boolean={}]'.format(str(value.value).lower()) # Transform the string into an assign for _help_ID if is_attr: target = Attribute(value=Name(id='self', ctx=Load()), attr=doc_id, ctx=Store()) else: target = Name(id=doc_id, ctx=Store()) help_str = s.value help_str.s = type_hint+s.value.s tree[n] = Assign(targets=[target], value=help_str) # Copy the line number from the original docstring copy_location(target, s) copy_location(tree[n], s) prev = s # Return the modified AST return tree
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 visit_Return(self, node): existing_node = self.generic_visit(node) value = existing_node.value if value is None: return existing_node return [Assign(targets=[Name(id=RESULT_NAME, ctx=Store())], value=value), self._create_context_call('return_value', [Name(id=RESULT_NAME, ctx=Load()), Num(n=existing_node.lineno)]), Return(value=Name(id=RESULT_NAME, ctx=Load()))]