Ejemplo n.º 1
0
 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])
Ejemplo n.º 2
0
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)