def test_underline_multiline():
doc = StmtBlockDoc([
ExprStmtDoc(IdDoc("foo")),
ExprStmtDoc(IdDoc("bar")),
])
doc.source_paths = [make_path("whole_doc")]
assert to_python_script(
doc, path_to_underline=make_path("whole_doc")) == format_script("""
foo
^^^
bar
^^^
""")
def get_func_doc_for_class(name):
args = [
AssignDoc(IdDoc("x"), rhs=None, annotation=IdDoc("int")),
AssignDoc(IdDoc("y"), rhs=LiteralDoc(1), annotation=IdDoc("int")),
]
body = [
AssignDoc(IdDoc("y"), OperationDoc(OperationKind.Add, [IdDoc("x"), LiteralDoc(1)])),
AssignDoc(IdDoc("y"), OperationDoc(OperationKind.Sub, [IdDoc("y"), LiteralDoc(1)])),
]
return FunctionDoc(
name=IdDoc(name),
args=args,
decorators=[IdDoc("wrap")],
return_type=LiteralDoc(None),
body=body,
)
def test_longer_prefix_must_win():
foo_x = IdDoc("foo_x")
foo_x.source_paths = [make_path("foo").attr("x")]
doc = StmtBlockDoc([
ExprStmtDoc(make_id_doc("foo")),
ExprStmtDoc(make_id_doc("bar")),
ExprStmtDoc(foo_x)
])
result = to_python_script(
doc, path_to_underline=make_path("foo").attr("x").attr("y"))
# "foo" should not be underlined because there is a document with a more specific path prefix
assert result == format_script("""
foo
bar
foo_x
^^^^^
""")
def test_print_class_doc(decorators, body, expected):
doc = ClassDoc(IdDoc("TestClass"), decorators, body)
assert to_python_script(doc) == format_script(expected)
def test_print_function_doc(args, decorators, body, return_type, expected):
doc = FunctionDoc(IdDoc("func"), args, decorators, return_type, body)
assert to_python_script(doc) == format_script(expected) # test
def test_print_expr_stmt_doc():
doc = ExprStmtDoc(CallDoc(IdDoc("f"), IdDoc("x")))
assert to_python_script(doc) == format_script("f(x)")
def test_print_scope_doc(lhs, body, expected):
doc = ScopeDoc(lhs, CallDoc(IdDoc("context")), body)
assert to_python_script(doc) == format_script(expected)
def test_print_id_doc(name):
doc = IdDoc(name)
assert to_python_script(doc) == format_script(name)
def test_print_if_doc(then_branch, else_branch, expected):
doc = IfDoc(IdDoc("pred"), then_branch, else_branch)
assert to_python_script(doc) == format_script(expected)
def test_print_slice_doc(slice_doc, expected):
doc = IdDoc("x")[slice_doc]
assert to_python_script(doc) == format_script(f"x[{expected}]")
def test_print_call_doc(args, kwargs, expected):
doc = CallDoc(IdDoc("f"), *args, **kwargs)
assert to_python_script(doc) == format_script(f"f{expected}")
def test_print_binary_operation_doc(op_kind, expected_token):
doc = OperationDoc(op_kind, [IdDoc("x"), IdDoc("y")])
assert to_python_script(doc) == format_script(f"x {expected_token} y")
def test_print_index_doc(indices, expected):
doc = IdDoc("x")[indices]
assert to_python_script(doc) == format_script(f"x{expected}")
def test_print_attr_doc(attr):
doc = IdDoc("x").attr(attr)
assert to_python_script(doc) == format_script(f"x.{attr}")
def generate_expr_precedence_test_cases():
x = IdDoc("x")
y = IdDoc("y")
z = IdDoc("z")
def negative(a):
return OperationDoc(OperationKind.USub, [a])
def invert(a):
return OperationDoc(OperationKind.Invert, [a])
def add(a, b):
return OperationDoc(OperationKind.Add, [a, b])
def sub(a, b):
return OperationDoc(OperationKind.Sub, [a, b])
def mult(a, b):
return OperationDoc(OperationKind.Mult, [a, b])
def div(a, b):
return OperationDoc(OperationKind.Div, [a, b])
def mod(a, b):
return OperationDoc(OperationKind.Mod, [a, b])
def pow(a, b):
return OperationDoc(OperationKind.Pow, [a, b])
def lshift(a, b):
return OperationDoc(OperationKind.LShift, [a, b])
def bit_and(a, b):
return OperationDoc(OperationKind.BitAnd, [a, b])
def bit_or(a, b):
return OperationDoc(OperationKind.BitOr, [a, b])
def bit_xor(a, b):
return OperationDoc(OperationKind.BitXor, [a, b])
def lt(a, b):
return OperationDoc(OperationKind.Lt, [a, b])
def eq(a, b):
return OperationDoc(OperationKind.Eq, [a, b])
def not_eq(a, b):
return OperationDoc(OperationKind.NotEq, [a, b])
def if_then_else(a, b, c):
return OperationDoc(OperationKind.IfThenElse, [a, b, c])
test_cases = {
"attr-call-index": [
(
add(x, y).attr("test"),
"(x + y).test",
),
(
add(x, y.attr("test")),
"x + y.test",
),
(
x[z].call(y),
"x[z](y)",
),
(
x.call(y)[z],
"x(y)[z]",
),
(
x.call(y).call(z),
"x(y)(z)",
),
(
x.call(y).attr("test"),
"x(y).test",
),
(
x.attr("test").call(y),
"x.test(y)",
),
(
x.attr("test").attr("test2"),
"x.test.test2",
),
(
LambdaDoc([x], x).call(y),
"(lambda x: x)(y)",
),
(
add(x, y)[z][add(z, z)].attr("name"),
"(x + y)[z][z + z].name",
),
],
"power": [
(
pow(pow(x, y), z),
"(x ** y) ** z",
),
(
pow(x, pow(y, z)),
"x ** y ** z",
),
(
pow(negative(x), negative(y)),
"(-x) ** -y",
),
(
pow(add(x, y), add(y, z)),
"(x + y) ** (y + z)",
),
],
"unary": [
(
invert(negative(y)),
"~-y",
),
(
negative(y).attr("test"),
"(-y).test",
),
(
negative(y.attr("test")),
"-y.test",
),
(
mult(negative(x), negative(y)),
"-x * -y",
),
(
negative(add(invert(x), negative(y))),
"-(~x + -y)",
),
],
"add-mult": [
(
mult(x, mult(y, z)),
"x * (y * z)",
),
(
mult(mult(x, y), z),
"x * y * z",
),
(
mult(x, add(y, z)),
"x * (y + z)",
),
(
mult(add(y, z), x),
"(y + z) * x",
),
(
add(x, mod(y, z)),
"x + y % z",
),
(
add(mult(y, z), x),
"y * z + x",
),
(
add(add(x, y), add(y, z)),
"x + y + (y + z)",
),
(
div(add(x, y), add(y, z)),
"(x + y) / (y + z)",
),
],
"shift": [
(
div(x, lshift(y, z)),
"x / (y << z)",
),
(
mult(lshift(y, z), x),
"(y << z) * x",
),
(
lshift(x, mult(y, z)),
"x << y * z",
),
(
lshift(mult(x, y), z),
"x * y << z",
),
(
lshift(mult(x, y), z),
"x * y << z",
),
(
lshift(lshift(x, y), z),
"x << y << z",
),
(
lshift(x, lshift(y, z)),
"x << (y << z)",
),
],
"bitwise": [
(
add(bit_or(x, y), bit_or(y, z)),
"(x | y) + (y | z)",
),
(
bit_and(bit_or(x, y), bit_or(y, z)),
"(x | y) & (y | z)",
),
(
bit_or(bit_and(x, y), bit_and(y, z)),
"x & y | y & z",
),
(
bit_and(bit_xor(x, bit_or(y, z)), z),
"(x ^ (y | z)) & z",
),
],
"comparison": [
(
not_eq(add(x, y), z),
"x + y != z",
),
(
eq(pow(x, y), z),
"x ** y == z",
),
(
lt(x, div(y, z)),
"x < y / z",
),
(
lt(x, if_then_else(y, y, y)),
"x < (y if y else y)",
),
],
"if-then-else": [
(
if_then_else(x, if_then_else(y, y, y), z),
"y if y else y if x else z",
),
(
if_then_else(if_then_else(x, x, x), y, z),
"y if (x if x else x) else z",
),
(
if_then_else(x, y, if_then_else(z, z, z)),
"y if x else (z if z else z)",
),
(
if_then_else(lt(x, x), add(y, y), mult(z, z)),
"y + y if x < x else z * z",
),
(
if_then_else(LambdaDoc([x], x), LambdaDoc([y], y), LambdaDoc([z], z)),
"(lambda y: y) if (lambda x: x) else (lambda z: z)",
),
],
"lambda": [
(
LambdaDoc([x, y], add(z, z)),
"lambda x, y: z + z",
),
(
add(LambdaDoc([x, y], z), z),
"(lambda x, y: z) + z",
),
(
LambdaDoc([x, y], add(z, z)).call(x, y),
"(lambda x, y: z + z)(x, y)",
),
(
LambdaDoc([x], LambdaDoc([y], z)),
"lambda x: lambda y: z",
),
],
}
return [
pytest.param(*args, id=f"{group_name}-{i}")
for group_name, cases in test_cases.items()
for i, args in enumerate(cases)
]
def printer(obj, object_path, ir_docsifier): # pylint: disable=unused-argument
return IdDoc(id_name)
def test_print_while_doc(body, expected):
doc = WhileDoc(IdDoc("pred"), body)
assert to_python_script(doc) == format_script(expected)
def test_print_for_doc(body, expected):
doc = ForDoc(IdDoc("x"), IdDoc("y"), body)
assert to_python_script(doc) == format_script(expected)
def make_id_doc(name: str, path_name: Optional[str] = None) -> IdDoc:
if path_name is None:
path_name = name
doc = IdDoc(name)
doc.source_paths = [make_path(path_name)]
return doc
assert to_python_script(doc) == format_script(f"x.{attr}")
@pytest.mark.parametrize(
"indices, expected",
[
(
(),
"[()]",
),
(
(LiteralDoc(1),),
"[1]",
),
(
(LiteralDoc(2), IdDoc("x")),
"[2, x]",
),
(
(SliceDoc(LiteralDoc(1), LiteralDoc(2)),),
"[1:2]",
),
(
(SliceDoc(LiteralDoc(1)), IdDoc("y")),
"[1:, y]",
),
(
(SliceDoc(), IdDoc("y")),
"[:, y]",
),
(