|
| 1 | +import ast |
| 2 | +from typing_extensions import assert_type |
| 3 | + |
| 4 | +# Test with source code strings |
| 5 | +assert_type(ast.parse("x = 1"), ast.Module) |
| 6 | +assert_type(ast.parse("x = 1", mode="exec"), ast.Module) |
| 7 | +assert_type(ast.parse("1 + 1", mode="eval"), ast.Expression) |
| 8 | +assert_type(ast.parse("x = 1", mode="single"), ast.Interactive) |
| 9 | +assert_type(ast.parse("(int, str) -> None", mode="func_type"), ast.FunctionType) |
| 10 | + |
| 11 | +# Test with mod objects - Module |
| 12 | +mod1: ast.Module = ast.Module([], []) |
| 13 | +assert_type(ast.parse(mod1), ast.Module) |
| 14 | +assert_type(ast.parse(mod1, mode="exec"), ast.Module) |
| 15 | +mod2: ast.Module = ast.Module(body=[ast.Expr(value=ast.Constant(value=42))], type_ignores=[]) |
| 16 | +assert_type(ast.parse(mod2), ast.Module) |
| 17 | + |
| 18 | +# Test with mod objects - Expression |
| 19 | +expr1: ast.Expression = ast.Expression(body=ast.Constant(value=42)) |
| 20 | +assert_type(ast.parse(expr1, mode="eval"), ast.Expression) |
| 21 | + |
| 22 | +# Test with mod objects - Interactive |
| 23 | +inter1: ast.Interactive = ast.Interactive(body=[]) |
| 24 | +assert_type(ast.parse(inter1, mode="single"), ast.Interactive) |
| 25 | + |
| 26 | +# Test with mod objects - FunctionType |
| 27 | +func1: ast.FunctionType = ast.FunctionType(argtypes=[], returns=ast.Constant(value=None)) |
| 28 | +assert_type(ast.parse(func1, mode="func_type"), ast.FunctionType) |
| 29 | + |
| 30 | +# Test that any AST node can be passed and returns the same type |
| 31 | +binop: ast.BinOp = ast.BinOp(left=ast.Constant(1), op=ast.Add(), right=ast.Constant(2)) |
| 32 | +assert_type(ast.parse(binop), ast.BinOp) |
| 33 | + |
| 34 | +constant: ast.Constant = ast.Constant(value=42) |
| 35 | +assert_type(ast.parse(constant), ast.Constant) |
| 36 | + |
| 37 | +expr_stmt: ast.Expr = ast.Expr(value=ast.Constant(value=42)) |
| 38 | +assert_type(ast.parse(expr_stmt), ast.Expr) |
| 39 | + |
| 40 | +# Test with additional parameters |
| 41 | +assert_type(ast.parse(mod1, filename="test.py"), ast.Module) |
| 42 | +assert_type(ast.parse(mod1, type_comments=True), ast.Module) |
| 43 | +assert_type(ast.parse(mod1, feature_version=(3, 10)), ast.Module) |
| 44 | +assert_type(ast.parse(binop, filename="test.py"), ast.BinOp) |
0 commit comments