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)