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)