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 bind_partial(self, *args, **kwds):
"""
Binds the provided positional and keyword arguments
and returns a new ``Function`` object with an updated signature.
"""
# We only need the signature, so clean the function body before eval'ing.
empty_func = self.replace(tree=replace_fields(self.tree, body=[ast.Pass()]))
signature = inspect.signature(empty_func.eval())
bargs = signature.bind_partial(*args, **kwds)
# Remove the bound arguments from the function AST
bound_argnames = set(bargs.arguments.keys())
new_tree = filter_function_def(self.tree, bound_argnames)
# Add assignments for bound parameters
assignments = []
gen_sym = GenSym.for_tree(new_tree)
new_bindings = {}
for name, value in bargs.arguments.items():
node, gen_sym, binding = reify_unwrapped(value, gen_sym)
new_bindings.update(binding)
assignments.append(ast.Assign(targets=[ast.Name(id=name, ctx=ast.Store())], value=node))
new_globals = dict(self.globals)
new_globals.update(new_bindings)
new_tree = replace_fields(new_tree, body=assignments + new_tree.body)
return Function(new_tree, new_globals, self.closure_vals, self._compiler_flags)
def fold(tree, constants):
statements = tree.body
cfg = build_cfg(statements)
gen_sym = GenSym.for_tree(tree)
new_nodes, temp_bindings = maximal_fixed_point(gen_sym, cfg.graph, cfg.enter, constants)
constants = dict(constants)
constants.update(temp_bindings)
new_tree = replace_exprs(tree, new_nodes)
return new_tree, constants
def _test_wrap_in_loop(body_src, expected_src, format_kwds={}, expected_bindings={}):
gen_sym = GenSym.for_tree()
body_nodes = ast.parse(unindent(body_src)).body
return_name = '_return_val'
gen_sym, inlined_body, new_bindings = _wrap_in_loop(gen_sym, body_nodes, return_name)
expected_body = ast.parse(unindent(expected_src.format(
return_val=return_name, **format_kwds))).body
assert_ast_equal(inlined_body, expected_body)
assert new_bindings == expected_bindings
def inline_functions(tree, constants):
gen_sym = GenSym.for_tree(tree)
constants = dict(constants)
state, tree = _inline_functions_walker(dict(gen_sym=gen_sym, constants=constants), tree)
return tree, state.constants