def test_mutiple_returns():
source = unindent('''
def f(x, y, z='foo'):
if x:
b = y + list(x)
return b
else:
return z
''')
tree = ast.parse(source)
expected_source = unindent('''
def f(__peval_mangled_1, __peval_mangled_2, __peval_mangled_3='foo'):
if __peval_mangled_1:
__peval_mangled_4 = __peval_mangled_2 + list(__peval_mangled_1)
return __peval_mangled_4
else:
return __peval_mangled_3
''')
expected_tree = ast.parse(expected_source)
gen_sym = GenSym.for_tree(tree)
gen_sym, new_tree = mangle(gen_sym, tree)
assert_ast_equal(new_tree, expected_tree)
def _inline(node, gen_sym, return_name, constants):
"""
Return a list of nodes, representing inlined function call.
"""
fn = constants[node.func.id]
fn_ast = Function.from_object(fn).tree
gen_sym, new_fn_ast = mangle(gen_sym, fn_ast)
parameter_assignments = _build_parameter_assignments(node, new_fn_ast)
body_nodes = new_fn_ast.body
gen_sym, inlined_body, new_bindings = _wrap_in_loop(gen_sym, body_nodes, return_name)
constants = dict(constants)
constants.update(new_bindings)
return parameter_assignments + inlined_body, gen_sym, constants