github-featured/apache--tvm
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
apache--tvm/tests/python/tvmscript/test_tvmscript_printer_python_doc_printer.py
T
Tianqi Chen 4fd6cfab15 [TVMScript] Render invisible paths in structural diagnostics (#19916) 
Structural diagnostics can identify a field below an object that
TVMScript renders without exposing that field. An underline alone then
points at the nearest visible parent and hides the full internal
location. This change keeps Script string-returning while making the
diagnostic context self-contained.

When the requested path is <root>.dtype, Script now returns this string
by default:

```text
Access path: <root>.dtype
Note: The underlined object is the nearest visible parent of this path.

T.int32
^^^^^^^
```

render_invisible_path_info defaults to true. Calls without target paths
are unchanged, and callers can set it to false to retain the legacy
underline-only string.

The implementation reuses the printer span-selection logic to capture
the deepest visible path and assembles the minimal
access-path/note/script block in C++. Pass-error enrichment uses the
same Script path. Production Python remains unchanged; focused Python
tests assert the complete strings for default-on, explicit-false,
hidden, exact-visible, unavailable-visible, pass-error, and
structural-equality cases.
2026-06-30 14:37:06 -04:00

1465 lines
35 KiB
Python
Raw Permalink Blame History

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
import itertools
import pytest
import tvm
from tvm.script.printer.doc import (
AssertDoc,
AssignDoc,
CallDoc,
ClassDoc,
CommentDoc,
DictDoc,
DocStringDoc,
ExprStmtDoc,
ForDoc,
FunctionDoc,
IdDoc,
IfDoc,
LambdaDoc,
ListDoc,
LiteralDoc,
OperationDoc,
OperationKind,
ReturnDoc,
ScopeDoc,
SliceDoc,
StmtBlockDoc,
TupleDoc,
WhileDoc,
)
from tvm.script.printer.doc_printer import to_python_script
def format_script(s: str) -> str:
"""
Remove leading and trailing blank lines, and make the minimum idention 0
"""
s = s.strip("\n")
non_empty_lines = [line for line in s.splitlines() if line and not line.isspace()]
if not non_empty_lines:
# no actual content
return ""
line_indents = [len(line) - len(line.lstrip(" ")) for line in non_empty_lines]
spaces_to_remove = min(line_indents)
cleaned_lines = "\n".join(line[spaces_to_remove:] for line in s.splitlines())
if not cleaned_lines.endswith("\n"):
cleaned_lines += "\n"
return cleaned_lines.strip()
@pytest.mark.parametrize(
"doc,expected",
[
(LiteralDoc(None), "None"),
(LiteralDoc(True), "True"),
(LiteralDoc(False), "False"),
(LiteralDoc("test"), '"test"'),
(LiteralDoc(""), '""'),
(LiteralDoc('""'), r'"\"\""'),
(LiteralDoc("\n\t\\test\r"), r'"\n\t\\test\r"'),
# TODO: fix the roundatrippable problem caused by utf8
pytest.param(LiteralDoc("\x88"), r'"\x88"', marks=pytest.mark.xfail),
(LiteralDoc(0), "0"),
(LiteralDoc(-1), "-1"),
(LiteralDoc(3.25), "3.25"),
(LiteralDoc(-0.5), "-0.5"),
],
ids=itertools.count(),
)
def test_print_literal_doc(doc, expected):
assert to_python_script(doc) == format_script(expected)
@pytest.mark.parametrize(
"name",
[
"test",
"_test",
"TestCase",
"test_case",
"test123",
],
ids=itertools.count(),
)
def test_print_id_doc(name):
doc = IdDoc(name)
assert to_python_script(doc) == format_script(name)
@pytest.mark.parametrize(
"attr",
[
"attr",
"_attr",
"Attr",
"attr_1",
],
ids=itertools.count(),
)
def test_print_attr_doc(attr):
doc = IdDoc("x").attr(attr)
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]",
),
(
(IdDoc("x"), IdDoc("y"), IdDoc("z")),
"[x, y, z]",
),
],
ids=itertools.count(),
)
def test_print_index_doc(indices, expected):
doc = IdDoc("x")[indices]
assert to_python_script(doc) == format_script(f"x{expected}")
UNARY_OP_TOKENS = {
OperationKind.USub: "-",
OperationKind.Invert: "~",
OperationKind.Not: "not ",
}
@pytest.mark.parametrize(
"op_kind, expected_token",
list(UNARY_OP_TOKENS.items()),
ids=UNARY_OP_TOKENS.keys(),
)
def test_print_unary_operation_doc(op_kind, expected_token):
doc = OperationDoc(op_kind, [IdDoc("x")])
assert to_python_script(doc) == format_script(f"{expected_token}x")
BINARY_OP_TOKENS = {
OperationKind.Add: "+",
OperationKind.Sub: "-",
OperationKind.Mult: "*",
OperationKind.Div: "/",
OperationKind.FloorDiv: "//",
OperationKind.Mod: "%",
OperationKind.Pow: "**",
OperationKind.LShift: "<<",
OperationKind.RShift: ">>",
OperationKind.BitAnd: "&",
OperationKind.BitOr: "|",
OperationKind.BitXor: "^",
OperationKind.Lt: "<",
OperationKind.LtE: "<=",
OperationKind.Eq: "==",
OperationKind.NotEq: "!=",
OperationKind.Gt: ">",
OperationKind.GtE: ">=",
OperationKind.And: "and",
OperationKind.Or: "or",
OperationKind.MatMul: "@",
}
@pytest.mark.parametrize(
"op_kind, expected_token",
list(BINARY_OP_TOKENS.items()),
ids=BINARY_OP_TOKENS.keys(),
)
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")
SPECIAL_OP_CASES = [
(
OperationKind.IfThenElse,
[LiteralDoc(True), LiteralDoc("true"), LiteralDoc("false")],
'"true" if True else "false"',
),
(
OperationKind.IfThenElse,
[IdDoc("x"), LiteralDoc(None), LiteralDoc(1)],
"None if x else 1",
),
]
@pytest.mark.parametrize(
"op_kind, operands, expected", SPECIAL_OP_CASES, ids=[kind for (kind, *_) in SPECIAL_OP_CASES]
)
def test_print_special_operation_doc(op_kind, operands, expected):
doc = OperationDoc(op_kind, operands)
assert to_python_script(doc) == format_script(expected)
def test_operation_doc_test_exhaustive():
special_op_covered = {k for k, *_ in SPECIAL_OP_CASES}
for op_kind in OperationKind:
if OperationKind._UnaryStart < op_kind < OperationKind._UnaryEnd:
assert op_kind in UNARY_OP_TOKENS, (
f"{op_kind.name} not covered in test_print_unary_operation_doc. "
f"Please add the expected token to UNARY_OP_TOKENS"
)
elif OperationKind._BinaryStart < op_kind < OperationKind._BinaryEnd:
assert op_kind in BINARY_OP_TOKENS, (
f"{op_kind.name} not covered in test_print_binary_operation_doc. "
f"Please add the expected token to BINARY_OP_TOKENS"
)
elif not op_kind.name.startswith("_"):
# Special Op
assert op_kind in special_op_covered, (
f"{op_kind.name} not covered in test_print_special_operation_doc. "
f"Please add the test cases for it to SPECIAL_OP_CASES"
)
@pytest.mark.parametrize(
"args, kwargs, expected",
[
(
(),
{},
"()",
),
(
(),
{"key0": IdDoc("u")},
"(key0=u)",
),
(
(),
{"key0": IdDoc("u"), "key1": IdDoc("v")},
"(key0=u, key1=v)",
),
(
(IdDoc("x"),),
{},
"(x)",
),
(
(IdDoc("x"),),
{"key0": IdDoc("u")},
"(x, key0=u)",
),
(
(IdDoc("x"),),
{"key0": IdDoc("u"), "key1": IdDoc("v")},
"(x, key0=u, key1=v)",
),
(
(IdDoc("x"), (IdDoc("y"))),
{},
"(x, y)",
),
(
(IdDoc("x"), (IdDoc("y"))),
{"key0": IdDoc("u")},
"(x, y, key0=u)",
),
(
(IdDoc("x"), (IdDoc("y"))),
{"key0": IdDoc("u"), "key1": IdDoc("v")},
"(x, y, key0=u, key1=v)",
),
],
ids=itertools.count(),
)
def test_print_call_doc(args, kwargs, expected):
doc = CallDoc(IdDoc("f"), *args, **kwargs)
assert to_python_script(doc) == format_script(f"f{expected}")
@pytest.mark.parametrize(
"args, expected",
[
(
(),
"lambda : 0",
),
(
(IdDoc("x"),),
"lambda x: 0",
),
(
(IdDoc("x"), IdDoc("y")),
"lambda x, y: 0",
),
(
(IdDoc("x"), IdDoc("y"), IdDoc("z")),
"lambda x, y, z: 0",
),
],
ids=itertools.count(),
)
def test_print_lambda_doc(args, expected):
doc = LambdaDoc(args, body=LiteralDoc(0))
assert to_python_script(doc) == format_script(expected)
@pytest.mark.parametrize(
"elements, expected",
[
(
(),
"[]",
),
(
[IdDoc("x")],
"[x]",
),
(
[IdDoc("x"), IdDoc("y")],
"[x, y]",
),
(
[IdDoc("x"), IdDoc("y"), IdDoc("z")],
"[x, y, z]",
),
],
ids=itertools.count(),
)
def test_print_list_doc(elements, expected):
doc = ListDoc(elements)
assert to_python_script(doc) == format_script(expected)
@pytest.mark.parametrize(
"elements, expected",
[
(
(),
"()",
),
(
[IdDoc("x")],
"(x,)",
),
(
[IdDoc("x"), IdDoc("y")],
"(x, y)",
),
(
[IdDoc("x"), IdDoc("y"), IdDoc("z")],
"(x, y, z)",
),
],
ids=itertools.count(),
)
def test_print_tuple_doc(elements, expected):
doc = TupleDoc(elements)
assert to_python_script(doc) == format_script(expected)
@pytest.mark.parametrize(
"content, expected",
[
(
{},
"{}",
),
(
{LiteralDoc("key_x"): IdDoc("x")},
'{"key_x": x}',
),
(
{LiteralDoc("key_x"): IdDoc("x"), LiteralDoc("key_y"): IdDoc("y")},
'{"key_x": x, "key_y": y}',
),
(
{
LiteralDoc("key_x"): IdDoc("x"),
LiteralDoc("key_y"): IdDoc("y"),
LiteralDoc("key_z"): IdDoc("z"),
},
'{"key_x": x, "key_y": y, "key_z": z}',
),
],
ids=itertools.count(),
)
def test_print_dict_doc(content, expected):
doc = DictDoc(content)
assert to_python_script(doc) == format_script(expected)
@pytest.mark.parametrize(
"slice_doc, expected",
[
(
SliceDoc(),
":",
),
(
SliceDoc(LiteralDoc(1)),
"1:",
),
(
SliceDoc(None, LiteralDoc(2)),
":2",
),
(
SliceDoc(LiteralDoc(1), LiteralDoc(2)),
"1:2",
),
(
SliceDoc(None, None, LiteralDoc(3)),
"::3",
),
(
SliceDoc(LiteralDoc(1), None, LiteralDoc(3)),
"1::3",
),
(
SliceDoc(None, LiteralDoc(2), LiteralDoc(3)),
":2:3",
),
(
SliceDoc(LiteralDoc(1), LiteralDoc(2), LiteralDoc(3)),
"1:2:3",
),
],
ids=itertools.count(),
)
def test_print_slice_doc(slice_doc, expected):
doc = IdDoc("x")[slice_doc]
assert to_python_script(doc) == format_script(f"x[{expected}]")
@pytest.mark.parametrize(
"stmts, expected",
[
(
[],
"",
),
(
[ExprStmtDoc(IdDoc("x"))],
"x",
),
(
[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],
"""
x
y
""",
),
],
ids=itertools.count(),
)
def test_print_stmt_block_doc(stmts, expected):
doc = StmtBlockDoc(stmts)
assert to_python_script(doc).strip() == format_script(expected).strip()
@pytest.mark.parametrize(
"doc, expected",
[
(
AssignDoc(IdDoc("x"), IdDoc("y"), None),
"x = y",
),
(
AssignDoc(IdDoc("x"), IdDoc("y"), IdDoc("int")),
"x: int = y",
),
(
AssignDoc(IdDoc("x"), None, IdDoc("int")),
"x: int",
),
(
AssignDoc(TupleDoc([IdDoc("x"), IdDoc("y")]), IdDoc("z"), None),
"x, y = z",
),
(
AssignDoc(TupleDoc([IdDoc("x"), TupleDoc([IdDoc("y"), IdDoc("z")])]), IdDoc("z"), None),
"x, (y, z) = z",
),
],
ids=itertools.count(),
)
def test_print_assign_doc(doc, expected):
assert to_python_script(doc) == format_script(expected)
@pytest.mark.parametrize(
"then_branch, else_branch, expected",
[
(
[ExprStmtDoc(IdDoc("x"))],
[],
"""
if pred:
x
""",
),
(
[],
[ExprStmtDoc(IdDoc("y"))],
"""
if pred:
pass
else:
y
""",
),
(
[ExprStmtDoc(IdDoc("x"))],
[ExprStmtDoc(IdDoc("y"))],
"""
if pred:
x
else:
y
""",
),
],
ids=itertools.count(),
)
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)
@pytest.mark.parametrize(
"body, expected",
[
(
[ExprStmtDoc(IdDoc("x"))],
"""
while pred:
x
""",
),
(
[],
"""
while pred:
pass
""",
),
],
ids=itertools.count(),
)
def test_print_while_doc(body, expected):
doc = WhileDoc(IdDoc("pred"), body)
assert to_python_script(doc) == format_script(expected)
@pytest.mark.parametrize(
"body, expected",
[
(
[ExprStmtDoc(IdDoc("x"))],
"""
for x in y:
x
""",
),
(
[],
"""
for x in y:
pass
""",
),
],
ids=itertools.count(),
)
def test_print_for_doc(body, expected):
doc = ForDoc(IdDoc("x"), IdDoc("y"), body)
assert to_python_script(doc) == format_script(expected)
@pytest.mark.parametrize(
"lhs, body, expected",
[
(
IdDoc("c"),
[ExprStmtDoc(IdDoc("x"))],
"""
with context() as c:
x
""",
),
(
IdDoc("c"),
[],
"""
with context() as c:
pass
""",
),
(
None,
[],
"""
with context():
pass
""",
),
(
None,
[ExprStmtDoc(IdDoc("x"))],
"""
with context():
x
""",
),
],
ids=itertools.count(),
)
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_expr_stmt_doc():
doc = ExprStmtDoc(CallDoc(IdDoc("f"), IdDoc("x")))
assert to_python_script(doc) == format_script("f(x)")
@pytest.mark.parametrize(
"msg, expected",
[
(
None,
"""
assert True
""",
),
(
LiteralDoc("test message"),
"""
assert True, "test message"
""",
),
],
ids=itertools.count(),
)
def test_print_assert_doc(msg, expected):
test = LiteralDoc(True)
doc = AssertDoc(test, msg)
assert to_python_script(doc) == format_script(expected)
@pytest.mark.parametrize(
"value, expected",
[
(
LiteralDoc(None),
"""
return None
""",
),
(
IdDoc("x"),
"""
return x
""",
),
],
ids=itertools.count(),
)
def test_print_return_doc(value, expected):
doc = ReturnDoc(value)
assert to_python_script(doc) == format_script(expected)
@pytest.mark.parametrize(
"args, decorators, return_type, body, expected",
[
(
[],
[],
None,
[],
"""
def func():
pass
""",
),
(
[AssignDoc(IdDoc("x"), rhs=None, annotation=IdDoc("int"))],
[],
IdDoc("int"),
[],
"""
def func(x: int) -> int:
pass
""",
),
(
[AssignDoc(IdDoc("x"), rhs=LiteralDoc(1), annotation=IdDoc("int"))],
[],
LiteralDoc(None),
[],
"""
def func(x: int = 1) -> None:
pass
""",
),
(
[],
[IdDoc("wrap")],
LiteralDoc(None),
[],
"""
@wrap
def func() -> None:
pass
""",
),
(
[],
[IdDoc("wrap_outter"), IdDoc("wrap_inner")],
LiteralDoc(None),
[],
"""
@wrap_outter
@wrap_inner
def func() -> None:
pass
""",
),
(
[
AssignDoc(IdDoc("x"), rhs=None, annotation=IdDoc("int")),
AssignDoc(IdDoc("y"), rhs=LiteralDoc(1), annotation=IdDoc("int")),
],
[IdDoc("wrap")],
LiteralDoc(None),
[],
"""
@wrap
def func(x: int, y: int = 1) -> None:
pass
""",
),
(
[
AssignDoc(IdDoc("x"), rhs=None, annotation=IdDoc("int")),
AssignDoc(IdDoc("y"), rhs=LiteralDoc(1), annotation=IdDoc("int")),
],
[IdDoc("wrap")],
LiteralDoc(None),
[
AssignDoc(IdDoc("y"), OperationDoc(OperationKind.Add, [IdDoc("x"), LiteralDoc(1)])),
AssignDoc(IdDoc("y"), OperationDoc(OperationKind.Sub, [IdDoc("y"), LiteralDoc(1)])),
],
"""
@wrap
def func(x: int, y: int = 1) -> None:
y = x + 1
y = y - 1
""",
),
],
ids=itertools.count(),
)
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 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,
)
@pytest.mark.parametrize(
"decorators, body, expected",
[
(
[],
[],
"""
class TestClass:
pass
""",
),
(
[IdDoc("wrap")],
[],
"""
@wrap
class TestClass:
pass
""",
),
(
[IdDoc("wrap_outter"), IdDoc("wrap_inner")],
[],
"""
@wrap_outter
@wrap_inner
class TestClass:
pass
""",
),
(
[IdDoc("wrap")],
[get_func_doc_for_class("f1")],
"""
@wrap
class TestClass:
@wrap
def f1(x: int, y: int = 1) -> None:
y = x + 1
y = y - 1
""",
),
(
[IdDoc("wrap")],
[get_func_doc_for_class("f1"), get_func_doc_for_class("f2")],
"""
@wrap
class TestClass:
@wrap
def f1(x: int, y: int = 1) -> None:
y = x + 1
y = y - 1
@wrap
def f2(x: int, y: int = 1) -> None:
y = x + 1
y = y - 1
""",
),
],
ids=itertools.count(),
)
def test_print_class_doc(decorators, body, expected):
doc = ClassDoc(IdDoc("TestClass"), decorators, body)
assert to_python_script(doc) == format_script(expected)
@pytest.mark.parametrize(
"comment, expected",
[
("", ""),
("test comment 1", "# test comment 1"),
(
"test comment 1\ntest comment 2",
"""
# test comment 1
# test comment 2
""",
),
],
ids=itertools.count(),
)
def test_print_comment_doc(comment, expected):
doc = CommentDoc(comment)
assert to_python_script(doc) == format_script(expected)
def test_print_comment_doc_preserves_trailing_tab():
assert to_python_script(CommentDoc("test comment\t")) == "# test comment\t"
@pytest.mark.parametrize(
"comment, expected",
[
(
"",
"",
),
(
"test comment 1",
'''
"""
test comment 1
"""
''',
),
(
"test comment 1\ntest comment 2",
'''
"""
test comment 1
test comment 2
"""
''',
),
],
ids=itertools.count(),
)
def test_print_doc_string_doc(comment, expected):
doc = DocStringDoc(comment)
assert to_python_script(doc) == format_script(expected)
@pytest.mark.parametrize(
"doc, comment, expected",
[
(
AssignDoc(IdDoc("x"), IdDoc("y"), IdDoc("int")),
"comment",
"""
x: int = y # comment
""",
),
(
IfDoc(IdDoc("x"), [ExprStmtDoc(IdDoc("y"))], [ExprStmtDoc(IdDoc("z"))]),
"comment",
"""
# comment
if x:
y
else:
z
""",
),
(
IfDoc(IdDoc("x"), [ExprStmtDoc(IdDoc("y"))], [ExprStmtDoc(IdDoc("z"))]),
"comment line 1\ncomment line 2",
"""
# comment line 1
# comment line 2
if x:
y
else:
z
""",
),
(
WhileDoc(
LiteralDoc(True),
[
AssignDoc(IdDoc("x"), IdDoc("y")),
],
),
"comment",
"""
# comment
while True:
x = y
""",
),
(
ForDoc(IdDoc("x"), IdDoc("y"), []),
"comment",
"""
# comment
for x in y:
pass
""",
),
(
ScopeDoc(IdDoc("x"), IdDoc("y"), []),
"comment",
"""
# comment
with y as x:
pass
""",
),
(
ExprStmtDoc(IdDoc("x")),
"comment",
"""
x # comment
""",
),
(
AssertDoc(LiteralDoc(True)),
"comment",
"""
assert True # comment
""",
),
(
ReturnDoc(LiteralDoc(1)),
"comment",
"""
return 1 # comment
""",
),
(
get_func_doc_for_class("f"),
"comment",
'''
@wrap
def f(x: int, y: int = 1) -> None:
"""
comment
"""
y = x + 1
y = y - 1
''',
),
(
get_func_doc_for_class("f"),
"comment line 1\n\ncomment line 3",
'''
@wrap
def f(x: int, y: int = 1) -> None:
"""
comment line 1
comment line 3
"""
y = x + 1
y = y - 1
''',
),
(
ClassDoc(IdDoc("TestClass"), decorators=[IdDoc("wrap")], body=[]),
"comment",
'''
@wrap
class TestClass:
"""
comment
"""
pass
''',
),
(
ClassDoc(IdDoc("TestClass"), decorators=[IdDoc("wrap")], body=[]),
"comment line 1\n\ncomment line 3",
'''
@wrap
class TestClass:
"""
comment line 1
comment line 3
"""
pass
''',
),
],
ids=itertools.count(),
)
def test_print_doc_comment(doc, comment, expected):
doc.comment = comment
assert to_python_script(doc) == format_script(expected)
@pytest.mark.parametrize(
"doc",
[
AssignDoc(IdDoc("x"), IdDoc("y"), IdDoc("int")),
ExprStmtDoc(IdDoc("x")),
AssertDoc(IdDoc("x")),
ReturnDoc(IdDoc("x")),
],
)
def test_print_invalid_multiline_doc_comment(doc):
doc.comment = "1\n2"
with pytest.raises(ValueError) as e:
to_python_script(doc)
assert "cannot have newline" in str(e.value)
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 not_(a):
return OperationDoc(OperationKind.Not, [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 and_(a, b):
return OperationDoc(OperationKind.And, [a, b])
def or_(a, b):
return OperationDoc(OperationKind.Or, [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)",
),
],
"boolean": [
(
not_(and_(x, y)),
"not (x and y)",
),
(
and_(not_(x), y),
"not x and y",
),
(
and_(or_(x, y), z),
"(x or y) and z",
),
(
or_(x, or_(y, z)),
"x or (y or z)",
),
(
or_(or_(x, y), z),
"x or y or z",
),
(
or_(and_(x, y), z),
# Maybe we should consider adding parentheses here
# for readability, even though it's not necessary.
"x and y or z",
),
(
and_(or_(not_(x), y), z),
"(not x or y) and z",
),
(
and_(lt(x, y), lt(y, z)),
"x < y and y < z",
),
(
or_(not_(eq(x, y)), lt(y, z)),
# Same as the previous one, the code here is not
# readable without parentheses.
"not x == y or y < z",
),
(
and_(if_then_else(x, y, z), x),
"(y if x else z) and x",
),
(
not_(if_then_else(x, y, z)),
"not (y if x else z)",
),
],
"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)
]
@pytest.mark.parametrize("doc, expected", generate_expr_precedence_test_cases())
def test_expr_precedence(doc, expected):
assert to_python_script(doc) == format_script(expected)
if __name__ == "__main__":
tvm.testing.main()
Reference in New Issue View Git Blame Copy Permalink