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])
def specialized_ast(fn_ast, global_bindings, *args, **kwargs):
''' Return AST of specialized function, and dict with closure bindings.
args and kwargs have the same meaning as in functools.partial.
Here we just handle the args and kwargs of function defenition.
'''
constants = dict(global_bindings)
fn_args = fn_ast.body[0].args
assert not fn_args.vararg and not fn_args.kwarg # TODO
if args:
for arg, value in zip(fn_args.args[:len(args)], args):
constants[arg.id] = value
del fn_args.args[:len(args)]
if kwargs:
arg_by_id = dict((arg.id, arg) for arg in fn_args.args)
for kwarg_name, kwarg_value in kwargs.iteritems():
constants[kwarg_name] = kwarg_value
fn_args.args.remove(arg_by_id[kwarg_name])
return optimized_ast(fn_ast, constants)
