def test_mutiple_returns(self):
source = '''
def f(x, y, z='foo'):
if x:
b = y + list(x)
return b
else:
return z
'''
ast_tree = ast.parse(shift_source(source))
expected_source = '''
def f(__ast_pe_var_4, __ast_pe_var_5, __ast_pe_var_6='foo'):
if __ast_pe_var_4:
__ast_pe_var_7 = __ast_pe_var_5 + list(__ast_pe_var_4)
__ast_pe_var_8 = __ast_pe_var_7
break
else:
__ast_pe_var_8 = __ast_pe_var_6
break
'''
inliner = Inliner(3, get_locals(ast_tree))
new_ast = inliner.visit(ast_tree)
self.assertASTEqual(new_ast, ast.parse(shift_source(expected_source)))
self.assertEqual(inliner.get_var_count(), 8)
self.assertEqual(inliner.get_return_var(), '__ast_pe_var_8')
self.assertEqual(inliner.get_bindings(), {
'x': '__ast_pe_var_4',
'y': '__ast_pe_var_5',
'z': '__ast_pe_var_6',
'b': '__ast_pe_var_7'})
def _test_opt(self, source, constants, expected_source=None, expected_new_bindings=None, print_source=False):
""" Test that with given constants, optimized_ast transforms
source to expected_source.
It :expected_new_bindings: is given, we check that they
are among new bindings returned by optimizer.
"""
if print_source:
print ast_to_string(ast.parse(shift_source(source)))
if expected_source is None:
expected_source = source
ast_tree = ast.parse(shift_source(source))
new_ast, bindings = optimized_ast(ast_tree, constants)
self.assertASTEqual(new_ast, ast.parse(shift_source(expected_source)))
if expected_new_bindings:
for k in expected_new_bindings:
if k not in bindings:
print "bindings:", bindings
self.assertEqual(bindings[k], expected_new_bindings[k])
