def test__bound_parameter_3():
func = """
a = 5
"""
module = ast.parse(func)
utils.set_parents(module)
assert module.body[0].targets[0].bound_parameter() == (None, None)
def test__set_parents():
module = ast.parse(inspect.getsource(utils))
utils.set_parents(module)
i = 0
for node in ast.walk(module):
assert hasattr(node, "parent")
i += 1
assert i > 2
def test__insert_badposition():
func = """
def foo():
...
"""
module = ast.parse(func)
utils.set_parents(module)
ellipsis = module.body[0].body[0]
stmt = ast.parse("a = 4").body[0]
with pytest.raises(utils.BadPosition):
result = ellipsis.insert(stmt, position="unknown")
def test__top_statement():
exprs = """
a = 1 + 2 + 3
3 + 4 + 5
b = 5 if True else 0
"""
module = ast.parse(exprs)
utils.set_parents(module)
expr1 = module.body[0].value.left.left
expr2 = module.body[1].value.left.left
expr3 = module.body[2].value.test
assert expr1.top_statement() is module.body[0]
assert expr2.top_statement() is module.body[1]
assert expr3.top_statement() is module.body[2]
def test__is_in_assign():
exprs = """
a = 1 + 2 + 3
3 + 4 + 5
b = 5 if True else 0
"""
module = ast.parse(exprs)
utils.set_parents(module)
expr1 = module.body[0].value.left.left
expr2 = module.body[1].value.left.left
expr3 = module.body[2].value.test
assert expr1.is_in_assign() is True
assert expr2.is_in_assign() is False
assert expr3.is_in_assign() is True
def test__next_assign():
func = """
a = 5
b = 4
a = a + b
"""
module = ast.parse(func)
utils.set_parents(module)
var_ref = module.body[0].targets[0]
assert var_ref.next_assign() is module.body[-1]
var_ref = module.body[1].targets[0]
assert var_ref.next_assign() is None
def test__bound_parameter():
func = """
def foo(a):
b = a + 5
c = 4
"""
module = ast.parse(func)
utils.set_parents(module)
foo = module.body[0]
variables = list(foo.all_variable_use())
assert variables[0].bound_parameter() == (None, None)
arg_a = foo.args.args[0]
assert variables[1].bound_parameter() == (foo, arg_a)
assert variables[2].bound_parameter() == (None, None)
def test__all_variable_use():
exprs = """
a + b + 5 + c if a in d else f
"""
module = ast.parse(exprs)
utils.set_parents(module)
expr = module.body[0].value
variables = list(expr.all_variable_use())
assert len(variables) == 6
assert variables[0].id == "a"
assert variables[1].id == "d"
assert variables[2].id == "a"
assert variables[3].id == "b"
assert variables[4].id == "c"
assert variables[5].id == "f"
def test__insert_expression_autoassign_off():
func = """
def foo():
return 3
"""
module = ast.parse(func)
utils.set_parents(module)
ret = module.body[0].body[0]
expr = ast.parse("1 + 3").body[0].value
result = ret.insert(expr, auto_assign=False)
assert isinstance(result, ast.Expr)
assert result is module.body[0].body[0]
assert result.value is expr
assert expr.parent is result
def test__insert_badtype():
func = """
def foo():
...
"""
module = ast.parse(func)
utils.set_parents(module)
ellipsis = module.body[0].body[0]
stmt = ast.parse("a = 4")
with pytest.raises(utils.BadASTNode) as exception_info:
result = ellipsis.insert(stmt, position="unknown")
exception = exception_info.value
assert exception.expected is ast.stmt
assert exception.actual is ast.Module
def test__all_variable_use():
func = """
def foo():
a = 4
b = 5
c = 6
a = 7
"""
module = ast.parse(func)
utils.set_parents(module)
foo = module.body[0]
variables = list(foo.all_variable_use())
assert "a" == variables[0].id
assert "b" == variables[1].id
assert "c" == variables[2].id
assert "a" == variables[3].id
assert len(variables) == 4
def test__insert_expression_autogenname():
func = """
def foo():
return 3
"""
module = ast.parse(func)
utils.set_parents(module)
ret = module.body[0].body[0]
expr = ast.parse("1 + 3").body[0].value
result = ret.insert(expr)
assert isinstance(result, ast.Assign)
assert result is module.body[0].body[0]
assert result.value is expr
assert expr.parent is result
assert len(result.targets) == 1
assert result.targets[0].id == f"@{id(expr)}"
def test__first_assign():
func = """
b = 4
a = 5
a = a + b
c
"""
module = ast.parse(func)
utils.set_parents(module)
var_ref = module.body[-2].value.left
assert var_ref.first_assignment() is module.body[1]
var_ref = module.body[-2].value.right
assert var_ref.first_assignment() is module.body[0]
var_ref = module.body[-1].value
assert var_ref.first_assignment() is None
def test__all_functions():
funcs = """
def f1():
def f2():
def f3():
...
def f4():
return 1
"""
module = ast.parse(funcs)
ellipsis = module.body[0].body[0].body[0].body[0]
utils.set_parents(module)
assert isinstance(ellipsis, ast.Expr)
assert isinstance(ellipsis.value, ast.Ellipsis)
all_funcs = tuple(f.name for f in ellipsis.all_parents_function())
assert ("f3", "f2", "f1") == all_funcs
assert "f4" not in all_funcs
def test__get_scopes():
funcs = """
def f1():
def f2():
def f3():
class A():
def f4():
...
def f5():
return 1
"""
module = ast.parse(funcs)
ellipsis = module.body[0].body[0].body[0].body[0].body[0].body[0].value
utils.set_parents(module)
assert isinstance(ellipsis, ast.Ellipsis)
scopes = tuple(getattr(f, "name", "module") for f in ellipsis.get_scopes())
assert ("f4", "A", "f3", "f2", "f1", "module") == scopes
assert "f5" not in scopes
def test__insert_before():
func = """
def foo():
return a
assert foo() == 4
"""
module = ast.parse(func)
utils.set_parents(module)
ret = module.body[0].body[0]
stmt = ast.parse("a = 4").body[0]
result = ret.insert(stmt)
assert stmt.parent is module.body[0]
assert result is stmt
assert result is module.body[0].body[0]
ast.fix_missing_locations(module)
obj = compile(module, filename="<ast>", mode="exec")
eval(obj)
def test__extract_autoinsert():
func = """
def foo():
return 1 + 2 + 1
assert foo() == 4
"""
module = ast.parse(func)
utils.set_parents(module)
ret = module.body[0].body[0]
ret_plus = module.body[0].body[0].value.left
assignement, var = ret_plus.extract(auto_insert=True)
assert assignement.parent is ret.parent
assert ret_plus.parent is assignement
assert var.parent is ret.value
ast.fix_missing_locations(module)
obj = compile(module, filename="<ast>", mode="exec")
eval(obj)
def test__unconditional_returns_return():
func = """
def foo():
if True:
return 1
else:
return 2
return 3
return 4
"""
module = ast.parse(func)
utils.set_parents(module)
foo = module.body[0]
returns = list(foo.unconditional_returns())
assert len(returns) == 2
assert returns[0].value.value == 3
assert returns[1].value.value == 4
assert foo.unconditional_return().value.value == 3
def test__insert_expression_controledname():
func = """
def foo():
return a
assert foo() == 4
"""
module = ast.parse(func)
utils.set_parents(module)
ret = module.body[0].body[0]
expr = ast.parse("1 + 3").body[0].value
result = ret.insert(expr, assign_name="a")
assert isinstance(result, ast.Assign)
assert result is module.body[0].body[0]
assert result.value is expr
assert expr.parent is result
ast.fix_missing_locations(module)
obj = compile(module, filename="<ast>", mode="exec")
eval(obj)
def test__extract_uncontrolled_name():
func = """
def foo():
return 1 + 2 + 1
assert foo() == 4
"""
module = ast.parse(func)
utils.set_parents(module)
ret = module.body[0].body[0]
ret_plus = module.body[0].body[0].value.left
assignement, var = ret_plus.extract()
assert assignement.parent is None
assert ret_plus.parent is assignement
assert var.parent is ret.value
ret.insert(assignement, position="before")
ast.fix_missing_locations(module)
obj = compile(module, filename="<ast>", mode="exec")
eval(obj)
def test__insert_after():
func = """
def foo():
a = 1
return a
assert foo() == 4
"""
module = ast.parse(func)
utils.set_parents(module)
assign = module.body[0].body[0]
stmt = ast.parse("a = 4").body[0]
result = assign.insert(stmt, position="after")
assert stmt.parent is module.body[0]
assert result is stmt
assert assign is module.body[0].body[0]
assert result is module.body[0].body[1]
ast.fix_missing_locations(module)
obj = compile(module, filename="<ast>", mode="exec")
eval(obj)
def test__bound_parameter_2():
func = """
def foo(a, b=4, *args, **kwargs):
b = a + 5
c = 4
args
kwargs
"""
module = ast.parse(func)
utils.set_parents(module)
foo = module.body[0]
variables = list(foo.all_variable_use())
arg_a = foo.args.args[0]
arg_b = foo.args.args[1]
args = foo.args.vararg
kwargs = foo.args.kwarg
assert variables[0].bound_parameter() == (foo, arg_b)
assert variables[1].bound_parameter() == (foo, arg_a)
assert variables[2].bound_parameter() == (None, None)
assert variables[3].bound_parameter() == (foo, args)
assert variables[4].bound_parameter() == (foo, kwargs)
