junrushao1994 commented on code in PR #12112:
URL: https://github.com/apache/tvm/pull/12112#discussion_r931777846
##########
tests/python/unittest/test_tvmscript_printer_python_doc_printer.py:
##########
@@ -421,7 +451,590 @@ def test_print_dict_doc(content, expected):
"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}]")
+
+
[email protected](
+ "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()
+
+
[email protected](
+ "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)
+
+
[email protected](
+ "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)
+
+
[email protected](
+ "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)
+
+
[email protected](
+ "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)
+
+
[email protected](
+ "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)")
+
+
[email protected](
+ "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)
+
+
[email protected](
+ "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)
+
+
[email protected](
+ "args, decorators, body, expected",
+ [
+ (
+ [],
+ [],
+ [],
+ """
+ def func() -> None:
+ pass
+ """,
+ ),
+ (
+ [AssignDoc(IdDoc("x"), rhs=None, annotation=IdDoc("int"))],
+ [],
+ [],
+ """
+ def func(x: int) -> None:
+ pass
+ """,
+ ),
+ (
+ [AssignDoc(IdDoc("x"), rhs=LiteralDoc(1),
annotation=IdDoc("int"))],
+ [],
+ [],
+ """
+ def func(x: int = 1) -> None:
+ pass
+ """,
+ ),
+ (
+ [],
+ [IdDoc("wrap")],
+ [],
+ """
+ @wrap
+ def func() -> None:
+ pass
+ """,
+ ),
+ (
+ [],
+ [IdDoc("wrap_outter"), IdDoc("wrap_inner")],
+ [],
+ """
+ @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")],
+ [],
+ """
+ @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")],
+ [
+ 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, expected):
+ doc = FunctionDoc(IdDoc("func"), args, decorators, LiteralDoc(None), 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,
+ )
+
+
[email protected](
+ "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)
+
+
[email protected](
+ "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)
+
+
[email protected](
+ "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)
Review Comment:
Let's append some epilogues to this file so that it could be runnable with
`python /path/to/file`
```python
if __name__ == "__main__":
tvm.testing.main()
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
