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 visit_Assign(self, node: ast.Assign):
if not isinstance(node.value, (ast.List, ast.Tuple)):
return self.generic_visit(node)
new_targets = []
for target in node.targets:
new_targets.append(self.visit(target))
node.targets = cast('List[ast.expr]', new_targets)
replacement_literal = ast.Call(func=ast.Name(self.literal_tracer,
ast.Load()),
args=[node.value],
keywords=[])
ast.copy_location(node.value, replacement_literal)
node.value = replacement_literal
return node
def visit_Assign(self, assign: ast.Assign) -> ast.Assign:
if isinstance(assign.value, ast.expr):
expr = ast.Expr(assign.value)
self.modify_expr(expr)
assign.value = expr.value
return assign
def visit_Assign(self, node: ast.Assign) -> ast.Assign:
self.found = False
self.generic_visit(node)
if self.found:
node.value = self.do_placeholder(node.value)
return node
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 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 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"])
def visit_Assign(self, node: ast.Assign):
new_targets = []
for target in node.targets:
new_targets.append(self.visit(target))
node.targets = new_targets
orig_value_id = id(node.value)
with fast.location_of(node.value):
node.value = self._make_tuple_event_for(
self.visit(node.value), TraceEvent.before_assign_rhs, orig_node_id=orig_value_id
)
node.value = fast.Call(
func=self._emitter_ast(),
args=[TraceEvent.after_assign_rhs.to_ast(), self._get_copy_id_ast(orig_value_id)],
keywords=fast.kwargs(ret=node.value),
)
return node
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))
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 visit_Assign(self, node: Assign) -> Any:
# if assigning to autogenerated variable
if len(node.targets) == 1 and node.targets[0].id.startswith( # type:ignore
self.context.prefix
):
node.value = self.replace_with_call(node.value) # type: ignore
return 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 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 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 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_Assign(self, node: ast.Assign) -> ast.Assign:
if len(node.targets) == 1:
target = node.targets[0]
if isinstance(target, ast.Name):
key = target.id
if key in self.mutations:
node.value = self.mutations[key]
self._used_mutations.remove(key)
return node
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 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 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 visit_Assign(self, assign: ast.Assign) -> ast.Assign:
assign.value = self.modify_expr(
assign.value,
return_stderr_and_returncode=isinstance(assign.targets[0],
ast.Tuple),
)
super().generic_visit(assign)
return assign
def getGlobalVariable(name): # _lambda_0
return Assign(
targets=[
Name(
id=name + '_locals',
ctx=Store()
)
],
value=Dict(keys=[], values=[]), 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)
