def test_custom_env():
@end_rewrite(file_name='test_custom_env.py')
@begin_rewrite(env=SymbolTable({'x': 1}, globals=globals()))
def f():
return x
assert f() == 1
def test_imbalanced(strict, a, b, c, d):
src = imbalanced_template.format(a, b, c, d)
tree = cst.parse_statement(src)
env = SymbolTable({}, {})
imbalanced = exec_def_in_file(tree, env)
can_name_error = False
for x in (False, True):
for y in (False, True):
try:
imbalanced(x, y)
except NameError:
can_name_error = True
break
if can_name_error and strict:
with pytest.raises(SyntaxError):
ssa_imbalanced = apply_passes([ssa(strict)])(imbalanced)
else:
ssa_imbalanced = apply_passes([ssa(strict)])(imbalanced)
for x in (False, True):
for y in (False, True):
try:
assert imbalanced(x, y) == ssa_imbalanced(x, y)
except NameError:
assert can_name_error
def test_call_in_annotations(strict, x, y):
r_x = x if x else 'int'
r_y = y if y else 'int'
x = f': {x}' if x else x
y = f': {y}' if y else y
src = sig_template.format(x, y, r_x, r_y)
tree = cst.parse_statement(src)
env = SymbolTable(locals(), globals())
f1 = exec_def_in_file(tree, env)
f2 = apply_passes([ssa(strict)])(f1)
def test_custom_env():
@apply_cst_passes(
[],
env=SymbolTable({'x': 1}, globals=globals()),
file_name='test_custom_env.py',
)
def f():
return x
assert f() == 1
def rewrite(self, tree: ast.AST, env: SymbolTable,
metadata: tp.MutableMapping) -> PASS_ARGS_T:
if not isinstance(self.phi, str):
phi_name = gen_free_name(tree, env, self.phi_name_prefix)
env.locals[phi_name] = self.phi
else:
phi_name = self.phi
visitor = IfExpTransformer(phi_name)
tree = visitor.visit(tree)
return tree, env, metadata
def test_basic_if(strict, a, b):
if a == b == 'return':
final = ''
else:
final = 'return r'
src = basic_template.format(a, b, final)
tree = cst.parse_statement(src)
env = SymbolTable({}, {})
basic = exec_def_in_file(tree, env)
ssa_basic = apply_passes([ssa(strict)])(basic)
for x in (False, True):
assert basic(x) == ssa_basic(x)
def test_basic_if(strict, a, b):
if a == b == 'return':
final = ''
else:
final = 'return r'
src = basic_template.format(a, b, final)
tree = ast.parse(src).body[0]
env = SymbolTable({}, {})
basic = exec_def_in_file(tree, env)
ssa_basic = _do_ssa(basic, strict, dump_src=True)
for x in (False, True):
assert basic(x) == ssa_basic(x)
def test_gen_free_name():
src = '''
class P:
P5 = 1
def __init__(self): self.y = 0
def P0():
return P.P5
P1 = P0()
'''
tree = ast.parse(src)
env = SymbolTable({}, {})
free_name = gen_free_name(tree, env)
assert free_name == '__auto_name_0'
free_name = gen_free_name(tree, env, prefix='P')
assert free_name == 'P2'
env = SymbolTable({'P3': 'foo'}, {})
free_name = gen_free_name(tree, env, prefix='P')
assert free_name == 'P2'
env = SymbolTable({'P3': 'foo'}, {'P2': 'bar'})
free_name = gen_free_name(tree, env, prefix='P')
assert free_name == 'P4'
def test_nested(strict, a, b, c, d):
if a == b == c == d == 'return':
final = ''
else:
final = 'return r'
src = nested_template.format(a, b, c, d, final)
tree = cst.parse_statement(src)
env = SymbolTable({}, {})
nested = exec_def_in_file(tree, env)
ssa_nested = apply_passes([ssa(strict)])(nested)
for x in (False, True):
for y in (False, True):
assert nested(x, y) == ssa_nested(x, y)
def test_nested(strict, a, b, c, d):
if a == b == c == d == 'return':
final = ''
else:
final = 'return r'
src = nested_template.format(a, b, c, d, final)
tree = ast.parse(src).body[0]
env = SymbolTable({}, {})
nested = exec_def_in_file(tree, env)
ssa_nested = _do_ssa(nested, strict, dump_src=True)
for x in (False, True):
for y in (False, True):
assert nested(x, y) == ssa_nested(x, y)
def gen_free_name(tree: ast.AST,
env: SymbolTable,
prefix: tp.Optional[str] = None) -> str:
names = used_names(tree) | env.keys()
if prefix is not None and prefix not in names:
return prefix
elif prefix is None:
prefix = '__auto_name_'
f_str = prefix + '{}'
c = 0
name = f_str.format(c)
while name in names:
c += 1
name = f_str.format(c)
return name
def test_gen_free_prefix():
src = '''
class P:
P5 = 1
def __init__(self): self.y = 0
def P0():
return P.P5
P1 = P0()
'''
tree = ast.parse(src)
env = SymbolTable({}, {})
free_prefix = gen_free_prefix(tree, env)
assert free_prefix == '__auto_prefix_0'
free_prefix = gen_free_prefix(tree, env, 'P')
assert free_prefix == 'P2'
def gen_free_prefix(tree: ast.AST,
env: SymbolTable,
preprefix: tp.Optional[str] = None) -> str:
def check_prefix(prefix: str, used_names: tp.AbstractSet[str]) -> bool:
return not any(name.startswith(prefix) for name in used_names)
names = used_names(tree) | env.keys()
if preprefix is not None and check_prefix(preprefix, names):
return preprefix
elif preprefix is None:
preprefix = '__auto_prefix_'
f_str = preprefix + '{}'
c = 0
prefix = f_str.format(c)
while not check_prefix(prefix, names):
c += 1
prefix = f_str.format(c)
return prefix
def test_imbalanced(strict, a, b, c, d):
src = imbalanced_template.format(a, b, c, d)
tree = ast.parse(src).body[0]
env = SymbolTable({}, {})
imbalanced = exec_def_in_file(tree, env)
can_name_error = False
for x in (False, True):
for y in (False, True):
try:
imbalanced(x, y)
except NameError:
can_name_error = True
break
if can_name_error and strict:
with pytest.raises(SyntaxError):
imbalanced_ssa = _do_ssa(imbalanced, strict, dump_src=True)
elif not can_name_error:
imbalanced_ssa = _do_ssa(imbalanced, strict, dump_src=True)
for x in (False, True):
for y in (False, True):
assert imbalanced(x, y) == imbalanced_ssa(x, y)
def is_free_prefix(tree: ast.AST, env: SymbolTable, prefix: str):
names = used_names(tree)
return not any(
name.startswith(prefix) for name in itertools.chain(names, env.keys()))
