def test_immutable_to_mutable(tree):
def _test(tree, mtree):
assert type(tree) is type(mtree)
if isinstance(tree, ast.AST):
for field, value in ast.iter_fields(tree):
_test(value, getattr(mtree, field))
elif isinstance(tree, list):
assert len(tree) == len(mtree)
for c, mc in zip(tree, mtree):
_test(c, mc)
else:
assert tree == mtree
itree = immutable_ast.immutable(tree)
mtree = immutable_ast.mutable(itree)
_test(tree, mtree)
def test_mutable_to_immutable(tree):
def _test(tree, itree):
if isinstance(tree, ast.AST):
assert isinstance(itree, immutable_ast.AST)
assert isinstance(tree, type(itree))
assert tree._fields == itree._fields
assert ImmutableMeta._mutable_to_immutable[type(tree)] is type(
itree)
for field, value in ast.iter_fields(tree):
_test(value, getattr(itree, field))
elif isinstance(tree, list):
assert isinstance(itree, tuple)
assert len(tree) == len(itree)
for c, ic in zip(tree, itree):
_test(c, ic)
else:
assert tree == itree
itree = immutable_ast.immutable(tree)
_test(tree, itree)
def test_eq(tree):
itree = immutable_ast.immutable(tree)
jtree = immutable_ast.immutable(tree)
assert itree == jtree
assert hash(itree) == hash(jtree)
def rewrite(self,
tree: ast.AST,
env: SymbolTable,
metadata: tp.MutableMapping) -> PASS_ARGS_T:
if not isinstance(tree, ast.FunctionDef):
raise TypeError('ssa should only be applied to functions')
# Going to use this in an assert later but need to get the info
# before any transformation happens
NR = _never_returns(tree.body)
# Find all attributes that are written
targets = collect_targets(tree, ast.Attribute)
replacer = AttrReplacer({})
init_reads = []
id_to_attr = {}
attr_to_name = {}
for t in targets:
i_t = immutable(t)
if not isinstance(t.value, ast.Name):
raise NotImplementedError(f'Only supports writing attributes '
f'of Name not {type(t.value)}')
elif i_t not in attr_to_name:
name = ast.Name(
id=gen_free_name(tree, env, '_'.join((t.value.id, t.attr))),
ctx=ast.Store())
# store the maping of names to attrs
id_to_attr[name.id] = t
attr_to_name[i_t] = name
#replace writes to the attr with writes to the name
replacer.add_replacement(t, name)
#replace reads to the attr with reads to the name
replacer.add_replacement(_flip_ctx(t), _flip_ctx(name))
# read the init value
if sys.version_info < (3, 8):
assign = ast.Assign(targets=[deepcopy(name)], value=_flip_ctx(t))
else:
assign = ast.Assign(
targets=[deepcopy(name)], value=_flip_ctx(t),
type_comment=None)
init_reads.append(immutable(assign))
else:
name = attr_to_name[i_t]
replacer.add_replacement(t, name)
replacer.add_replacement(_flip_ctx(t), _flip_ctx(name))
# Replace references to the attr with the name generated above
tree = replacer.visit(tree)
# insert initial reads
tree.body = [mutable(r) for r in init_reads] + tree.body
# Perform ssa
r_name = gen_free_prefix(tree, env, self.return_prefix)
visitor = SSATransformer(env, r_name, id_to_attr.keys(), self.strict)
tree = visitor.visit(tree)
#insert the write backs to the attrs
for name, conditons in visitor.attr_states.items():
if conditons:
tree.body.append(
ast.Assign(
targets=[deepcopy(id_to_attr[name])],
value=_fold_conditions(conditons)
)
)
else:
tree.body.append(
ast.Assign(
targets=[deepcopy(id_to_attr[name])],
value=ast.Name(visitor.name_table[name], ast.Load())
)
)
# insert the return
if visitor.returns:
tree.body.append(
ast.Return(
value=_fold_conditions(visitor.returns)
)
)
else:
assert NR
return tree, env, metadata
def vist_Call(self, node: ast.Call):
if self.count_calls:
self.cses[immutable(node)] += 1
return self.generic_visit(node)
def _get_key(self, node):
if isinstance(node, ast.Attribute):
# Need the immutable value so its comparable
return immutable(node)
else:
return None
def visit_IfExpr(self, node: ast.IfExp):
self.cses[immutable(node)] += 1
self.generic_visit(node)
def visit_Compare(self, node: ast.Compare):
self.cses[immutable(node)] += 1
self.generic_visit(node)
def visit_BoolOp(self, node: ast.BoolOp):
self.cses[immutable(node)] += 1
self.generic_visit(node)
def visit_UnaryOp(self, node: ast.UnaryOp):
self.cses[immutable(node)] += 1
self.generic_visit(node)
def _get_key(node: ast.AST):
if isinstance(node, ast.expr):
# Need the immutable value so its comparable
return immutable(node)
else:
return None
