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