Ejemplo n.º 1
0
def test_reify_unwrapped():
    class Dummy(): pass
    x = Dummy()
    gen_sym = GenSym()
    node, gen_sym, binding = reify_unwrapped(x, gen_sym)
    assert_ast_equal(node, ast.Name(id='__peval_temp_1', ctx=ast.Load()))
    assert binding == dict(__peval_temp_1=x)
Ejemplo n.º 2
0
def check_partial_apply(func, args=None, kwds=None,
        expected_source=None, expected_new_bindings=None):
    ''' 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 args is None:
        args = tuple()
    if kwds is None:
        kwds = {}

    new_func = partial_apply(func, *args, **kwds)
    function = Function.from_object(new_func)

    if expected_source is not None:
        assert_ast_equal(function.tree, ast.parse(unindent(expected_source)).body[0])

    if expected_new_bindings is not None:
        for k in expected_new_bindings:
            if k not in function.globals:
                print('Expected binding missing:', k)

            binding = function.globals[k]
            expected_binding = expected_new_bindings[k]

            # Python 3.2 defines equality for range objects incorrectly
            # (namely, the result is always False).
            # So we just test it manually.
            if sys.version_info < (3, 3) and isinstance(expected_binding, range):
                assert type(binding) == type(expected_binding)
                assert list(binding) == list(expected_binding)
            else:
                assert binding == expected_binding
Ejemplo n.º 3
0
def test_visit_after():

    @ast_transformer
    def simplify(node, visit_after, visiting_after, **kwds):
        if isinstance(node, ast.If):
            if not visiting_after:
                visit_after()
                return node

            # This wouldn't work if we didn't simplify the child nodes first
            if (len(node.orelse) == 0 and len(node.body) == 1
                    and isinstance(node.body[0], ast.Pass)):
                return ast.Pass()
            else:
                return node
        else:
            return node

    node = get_ast(dummy_if)
    new_node = simplify(node)

    assert_ast_equal(new_node, get_ast(
        """
        def dummy_if():
            pass
        """))
Ejemplo n.º 4
0
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)
Ejemplo n.º 5
0
def test_value_to_node():
    class Dummy(): pass
    x = Dummy()
    gen_sym = GenSym()
    node, gen_sym, binding = value_to_node(x, gen_sym)
    assert_ast_equal(node, ast.Name(id='__peval_temp_1', ctx=ast.Load()))
    assert binding == dict(__peval_temp_1=x)
Ejemplo n.º 6
0
def check_peval_expression(source, bindings, expected_source,
        fully_evaluated=False, expected_value=None, expected_temp_bindings=None):

    source_tree = expression_ast(source)

    # In some cases we need to enforce the expected node,
    # because it cannot be obtained by parsing
    # (e.g. "-5" is parsed as "UnaryOp(op=USub(), Num(n=5))", not as "Num(n=-5)").
    # But we expect the latter from a fully evaluated expression.
    if isinstance(expected_source, str):
        expected_tree = expression_ast(expected_source)
    else:
        expected_tree = expected_source

    gen_sym = GenSym()
    result, gen_sym = peval_expression(source_tree, gen_sym, bindings)

    assert_ast_equal(result.node, expected_tree)

    assert result.fully_evaluated == fully_evaluated
    if fully_evaluated:
        assert result.value == expected_value

    if expected_temp_bindings is not None:
        for key, val in expected_temp_bindings.items():
            assert key in result.temp_bindings
            assert result.temp_bindings[key] == expected_temp_bindings[key]
Ejemplo n.º 7
0
def check_partial_apply(func, args=None, kwds=None,
        expected_source=None, expected_new_bindings=None):
    """
    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 args is None:
        args = tuple()
    if kwds is None:
        kwds = {}

    new_func = partial_apply(func, *args, **kwds)
    function = Function.from_object(new_func)

    if expected_source is not None:
        assert_ast_equal(function.tree, ast.parse(unindent(expected_source)).body[0])

    if expected_new_bindings is not None:
        for k in expected_new_bindings:
            if k not in function.globals:
                print('Expected binding missing:', k)

            binding = function.globals[k]
            expected_binding = expected_new_bindings[k]

            assert binding == expected_binding
Ejemplo n.º 8
0
def check_reify(value, expected_ast, preferred_name=None, expected_binding=None):
    kvalue = KnownValue(value, preferred_name=preferred_name)
    gen_sym = GenSym()
    node, gen_sym, binding = reify(kvalue, gen_sym)

    assert_ast_equal(node, expected_ast)
    if expected_binding is not None:
        assert binding == expected_binding
Ejemplo n.º 9
0
def check_node_to_maybe_kvalue(node, bindings, expected_result, expected_preferred_name=None):
    node_or_kvalue = node_to_maybe_kvalue(node, bindings)

    if is_known_value(node_or_kvalue):
        assert node_or_kvalue.value == expected_result
        assert node_or_kvalue.preferred_name == expected_preferred_name
    else:
        assert_ast_equal(node_or_kvalue, expected_result)
Ejemplo n.º 10
0
def _test_build_parameter_assignments(call_str, signature_str, expected_assignments):

    call_node = ast.parse("func(" + call_str + ")").body[0].value
    signature_node = ast.parse("def func(" + signature_str + "):\n\tpass").body[0]

    assignments = _build_parameter_assignments(call_node, signature_node)

    expected_assignments = ast.parse(unindent(expected_assignments)).body

    assert_ast_equal(assignments, expected_assignments)
Ejemplo n.º 11
0
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
Ejemplo n.º 12
0
def test_change_node():

    @ast_transformer
    def change_name(node, **kwds):
        if isinstance(node, ast.Name) and node.id == 'a':
            return ast.Name(id='b', ctx=node.ctx)
        else:
            return node

    node = get_ast(dummy)
    new_node = check_mutation(node, change_name)
    assert_ast_equal(new_node, get_ast("""
        def dummy(x, y):
            c = 4
            b = 1
        """))
Ejemplo n.º 13
0
def _test_replace_returns(source, expected_source, expected_returns_ctr, expected_returns_in_loops):

    nodes = ast.parse(unindent(source)).body

    return_var = 'return_var'
    return_flag_var = 'return_flag'

    expected_source = expected_source.format(
        return_var=return_var, return_flag=return_flag_var)
    expected_nodes = ast.parse(unindent(expected_source)).body

    new_nodes, returns_ctr, returns_in_loops = _replace_returns(
        nodes, return_var, return_flag_var)

    assert_ast_equal(new_nodes, expected_nodes)
    assert returns_ctr == expected_returns_ctr
    assert returns_in_loops == expected_returns_in_loops
Ejemplo n.º 14
0
def test_walker():

    @ast_walker
    def process_numbers(node, state, **kwds):
        if isinstance(node, ast.Num):
            return ast.Num(n=node.n + 1), state.update(numbers=state.numbers.add(node.n))
        else:
            return node, state

    node = get_ast(dummy)
    new_node, state = process_numbers(node, state=dict(numbers=immutableset()))

    assert state.numbers == set([1, 4])
    assert_ast_equal(new_node, get_ast("""
        def dummy(x, y):
            c = 5
            a = 2
        """))
Ejemplo n.º 15
0
def test_global_context():

    @ast_transformer
    def rename(node, ctx, **kwds):
        if isinstance(node, ast.Name) and node.id == ctx.old_name:
            return ast.Name(id=ctx.new_name, ctx=node.ctx)
        else:
            return node

    node = get_ast(dummy)
    new_node = rename(node, ctx=dict(old_name='c', new_name='d'))

    assert_ast_equal(new_node, get_ast(
        """
        def dummy(x, y):
            d = 4
            a = 1
        """))
Ejemplo n.º 16
0
def test_add_statement():

    @ast_transformer
    def add_statement(node, **kwds):
        if isinstance(node, ast.Assign):
            return [node, ast.parse("b = 2").body[0]]
        else:
            return node

    node = get_ast(dummy)
    new_node = check_mutation(node, add_statement)
    assert_ast_equal(new_node, get_ast("""
        def dummy(x, y):
            c = 4
            b = 2
            a = 1
            b = 2
        """))
Ejemplo n.º 17
0
def test_list_element():
    """
    Tests the removal of an AST node that is an element of a list
    referenced by a field of the parent node.
    """

    @ast_transformer
    def remove_list_element(node, **kwds):
        if isinstance(node, ast.Assign) and node.targets[0].id == 'a':
            return None
        else:
            return node

    node = get_ast(dummy)
    new_node = check_mutation(node, remove_list_element)
    assert_ast_equal(new_node, get_ast("""
        def dummy(x, y):
            c = 4
        """))
Ejemplo n.º 18
0
def test_remove_field():
    """
    Tests the removal of an AST node that is referenced by a field of the parent node.
    """

    @ast_transformer
    def remove_field(node, **kwds):
        if isinstance(node, ast.arg) and node.arg == 'x':
            return None
        else:
            return node

    node = get_ast(dummy)
    new_node = check_mutation(node, remove_field)
    assert_ast_equal(new_node, get_ast("""
        def dummy(y):
            c = 4
            a = 1
        """))
Ejemplo n.º 19
0
def _test_replace_returns(source, expected_source, expected_returns_ctr, expected_returns_in_loops):

    nodes = ast.parse(unindent(source)).body

    true_val = 'true_val'
    return_var = 'return_var'
    return_flag_var = 'return_flag'

    expected_source = expected_source.format(
        return_var=return_var, return_flag=return_flag_var, true_val='true_val')
    expected_nodes = ast.parse(unindent(expected_source)).body

    true_node = ast.Name(true_val, ast.Load())
    new_nodes, returns_ctr, returns_in_loops = _replace_returns(
        nodes, return_var, return_flag_var, true_node)

    assert_ast_equal(new_nodes, expected_nodes)
    assert returns_ctr == expected_returns_ctr
    assert returns_in_loops == expected_returns_in_loops
Ejemplo n.º 20
0
def test_walk_field_transform():

    @ast_transformer
    def increment(node, walk_field, **kwds):
        if isinstance(node, ast.Assign):
            return replace_fields(node, targets=node.targets, value=walk_field(node.value))
        elif isinstance(node, ast.Num):
            return ast.Num(n=node.n + 1)
        else:
            return node

    node = get_ast(dummy)
    new_node = increment(node)

    assert_ast_equal(new_node, get_ast(
        """
        def dummy(x, y):
            c = 5
            a = 2
        """))
Ejemplo n.º 21
0
def test_skip_fields():

    @ast_transformer
    def increment(node, skip_fields, **kwds):
        if isinstance(node, ast.Assign) and node.targets[0].id == 'c':
            skip_fields()

        if isinstance(node, ast.Num):
            return ast.Num(n=node.n + 1)
        else:
            return node

    node = get_ast(dummy)
    new_node = increment(node)

    assert_ast_equal(new_node, get_ast(
        """
        def dummy(x, y):
            c = 4
            a = 2
        """))
Ejemplo n.º 22
0
def test_block_autofix():

    # This transformer removes If nodes from statement blocks,
    # but it has no way to check whether the resulting body still has some nodes or not.
    # That's why the walker adds a Pass node automatically if after all the transformations
    # a statement block turns out to be empty.
    @ast_transformer
    def delete_ifs(node, **kwds):
        if isinstance(node, ast.If):
            return None
        else:
            return node

    node = get_ast(dummy_if)
    new_node = delete_ifs(node)

    assert_ast_equal(new_node, get_ast(
        """
        def dummy_if():
            pass
        """))
Ejemplo n.º 23
0
def test_walk_field_transform_inspect():

    @ast_walker
    def names_and_incremented_nums(node, state, walk_field, **kwds):
        if isinstance(node, ast.Assign):
            value_node, state = walk_field(node.value, state)
            new_node = replace_fields(node, targets=node.targets, value=value_node)
            new_state = state.update(objs=state.objs.add(node.targets[0].id))
            return new_node, new_state
        elif isinstance(node, ast.Num):
            return ast.Num(n=node.n + 1), state.update(objs=state.objs.add(node.n))
        else:
            return node, state

    node = get_ast(dummy)
    new_node, state = names_and_incremented_nums(node, state=dict(objs=immutableset()))
    assert state.objs == set(['a', 'c', 1, 4])
    assert_ast_equal(new_node, get_ast(
        """
        def dummy(x, y):
            c = 5
            a = 2
        """))
Ejemplo n.º 24
0
def test_generator_exp():

    # Need to do this manually, since we can't compare generator expressions
    # without changing their state.

    source_tree = expression_ast("(x + 1 for x in range(a))")
    expected_tree = expression_ast("__peval_temp_2")
    bindings = dict(a=10, range=range)

    gen_sym = GenSym()
    result, gen_sym = peval_expression(source_tree, gen_sym, bindings)

    assert_ast_equal(result.node, expected_tree)
    assert result.fully_evaluated

    expected_genexp = (x + 1 for x in range(10))

    assert type(result.value) == type(expected_genexp)
    assert list(result.value) == list(expected_genexp)

    # Since the binding contained the reference to the same genexp,
    # and we destroyed it, we need to do the evaluation again
    # in order to check the binding as well

    gen_sym = GenSym()
    result, gen_sym = peval_expression(source_tree, gen_sym, bindings)

    expected_genexp = (x + 1 for x in range(10))

    assert '__peval_temp_2' in result.temp_bindings
    binding = result.temp_bindings['__peval_temp_2']
    assert type(binding) == type(expected_genexp)
    assert list(binding) == list(expected_genexp)


    check_peval_expression(
        '(x + 1 for x in range(a))', dict(a=10), '(x + 1 for x in range(10))')
Ejemplo n.º 25
0
def test_walk_children():

    @ast_transformer
    def mangle_outer_functions(node, **kwds):
        if isinstance(node, ast.FunctionDef):
            return replace_fields(node, name='__' + node.name)
        else:
            return node

    @ast_transformer
    def mangle_all_functions(node, walk_field, **kwds):
        if isinstance(node, ast.FunctionDef):
            return replace_fields(
                node, name='__' + node.name, body=walk_field(node.body, block_context=True))
        else:
            return node

    node = get_ast(dummy_nested)

    new_node = mangle_outer_functions(node)
    assert_ast_equal(new_node, get_ast(
        """
        def __dummy_nested(x, y):
            def inner_function(z):
                return z
            return inner_function
        """))

    new_node = mangle_all_functions(node)
    assert_ast_equal(new_node, get_ast(
        """
        def __dummy_nested(x, y):
            def __inner_function(z):
                return z
            return inner_function
        """))
Ejemplo n.º 26
0
def test_prepend():

    @ast_transformer
    def prepender(node, prepend, **kwds):
        if isinstance(node, ast.Name):
            if node.id == 'a':
                prepend(
                    [ast.Assign(targets=[ast.Name(id='k', ctx=ast.Store())], value=ast.Num(n=10))])
                return node
            elif node.id == 'b':
                prepend(
                    [ast.Assign(targets=[ast.Name(id='l', ctx=ast.Store())], value=ast.Num(n=20))])
                return ast.Name(id='d', ctx=node.ctx)
            elif node.id == 'c':
                prepend(
                    [ast.Assign(targets=[ast.Name(id='m', ctx=ast.Store())], value=ast.Num(n=30))])
                return node
            else:
                return node
        else:
            return node

    node = get_ast(dummy_blocks)
    new_node = prepender(node)

    assert_ast_equal(new_node, get_ast(
        """
        def dummy_blocks(x, y):
            k = 10
            a = 1
            if x:
                l = 20
                d = 2
            m = 30
            c = 3
        """))
Ejemplo n.º 27
0
def check_mutation(node, walker):
    node_ref = copy.deepcopy(node)
    new_node = walker(node)
    assert ast.dump(node) != ast.dump(new_node)
    assert_ast_equal(node, node_ref)
    return new_node