# 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.

# ruff: noqa: F841

"""

In this test file, we want to make sure the Python code can construct

Doc objects, then access and modify their attributes correctly.

"""

import pytest

from tvm_ffi.access_path import AccessPath

import tvm

from tvm.script.printer.doc import (

AssertDoc,

AssignDoc,

AttrAccessDoc,

CallDoc,

ClassDoc,

CommentDoc,

DictDoc,

DocStringDoc,

ExprStmtDoc,

ForDoc,

FunctionDoc,

IdDoc,

IfDoc,

IndexDoc,

LambdaDoc,

ListDoc,

LiteralDoc,

OperationDoc,

OperationKind,

ReturnDoc,

ScopeDoc,

SliceDoc,

StmtBlockDoc,

TupleDoc,

WhileDoc,

)

@pytest.mark.parametrize(

"value",

[None, "test", 0, 1, -2, 0.0, 1.5, -1.3, True, False],

)

def test_literal_doc_construction(value):

doc = LiteralDoc(value)

if isinstance(value, float):

# FloatImm cannot be compared with Python's float directly

assert float(doc.value) == pytest.approx(value)

else:

assert doc.value == value

def test_id_doc():

doc = IdDoc("name")

assert doc.name == "name"

def test_attr_access_doc():

target = IdDoc("x")

doc = AttrAccessDoc(target, "attribute")

assert doc.value == target

assert doc.name == "attribute"

@pytest.mark.parametrize(

"indices",

[

[],

[LiteralDoc(1)],

[LiteralDoc(2), IdDoc("x")],

[SliceDoc(LiteralDoc(1), LiteralDoc(2))],

[SliceDoc(LiteralDoc(1)), IdDoc("y")],

],

)

def test_index_doc(indices):

target = IdDoc("x")

doc = IndexDoc(target, indices)

assert doc.value == target

assert list(doc.indices) == indices

@pytest.mark.parametrize(

"args, kwargs",

[

([], {}),

([LiteralDoc("arg")], {}),

([LiteralDoc("arg"), IdDoc("x")], {}),

([], {"x": LiteralDoc("x")}),

([], {"x": LiteralDoc("x"), "y": LiteralDoc("y")}),

([LiteralDoc("arg")], {"x": LiteralDoc("x"), "y": LiteralDoc("y")}),

([LiteralDoc("arg"), IdDoc("x")], {"x": LiteralDoc("x"), "y": LiteralDoc("y")}),

],

)

def test_call_doc(args, kwargs):

target = IdDoc("x")

doc = CallDoc(target, *args, **kwargs)

assert doc.callee == target

assert list(doc.args) == args

assert dict(zip(doc.kwargs_keys, doc.kwargs_values)) == kwargs

@pytest.mark.parametrize(

"operands",

[

[],

[LiteralDoc(1)],

[LiteralDoc(2), IdDoc("x")],

[LiteralDoc(2), IdDoc("x"), LiteralDoc("y")],

],

)

def test_operation_doc(operands):

# Here we just test the contructor and attr visitor of OperationDoc

# so the choice of OperationKind doesn't matter

operator = OperationKind.Add

doc = OperationDoc(OperationKind.Add, operands)

assert doc.kind == operator

assert list(doc.operands) == operands

@pytest.mark.parametrize(

"args",

[

[],

[IdDoc("x")],

[IdDoc("x"), IdDoc("y")],

],

)

def test_lambda_doc(args):

body = LiteralDoc(1)

doc = LambdaDoc(args, body)

assert doc.body == body

assert list(doc.args) == args

@pytest.mark.parametrize(

"elements",

[

[],

[IdDoc("x")],

[IdDoc("x"), IdDoc("y")],

],

)

def test_tuple_doc(elements):

doc = TupleDoc(elements)

assert list(doc.elements) == elements

@pytest.mark.parametrize(

"elements",

[

[],

[IdDoc("x")],

[IdDoc("x"), IdDoc("y")],

],

)

def test_list_doc(elements):

doc = ListDoc(elements)

assert list(doc.elements) == elements

@pytest.mark.parametrize(

"content",

[

{},

{LiteralDoc("k"): IdDoc("v")},

{LiteralDoc("k"): IdDoc("v"), LiteralDoc("k2"): IdDoc("v2")},

],

)

def test_dict_doc(content):

doc = DictDoc(content)

assert dict(zip(doc.keys, doc.values)) == content

@pytest.mark.parametrize("start", [LiteralDoc(1), None])

@pytest.mark.parametrize("stop", [LiteralDoc(2), None])

@pytest.mark.parametrize("step", [LiteralDoc(3), None])

def test_slice_doc(start, stop, step):

doc = SliceDoc(start, stop)

assert doc.start == start

assert doc.stop == stop

def test_expr_doc_attr_access():

target = IdDoc("x")

attr = "test"

doc = target.attr(attr)

assert doc.value == target

assert doc.name == attr

@pytest.mark.parametrize(

"indices",

[

(),

LiteralDoc(1),

SliceDoc(LiteralDoc(1), LiteralDoc(2)),

(LiteralDoc(1),),

(LiteralDoc(2), IdDoc("x")),

(SliceDoc(LiteralDoc(1), LiteralDoc(2)),),

(SliceDoc(LiteralDoc(1)), IdDoc("y")),

],

)

def test_expr_doc_get_item(indices):

target = IdDoc("x")

doc = target[indices]

assert doc.value == target

if not isinstance(indices, tuple):

indices = (indices,)

assert tuple(doc.indices) == indices

@pytest.mark.parametrize(

"args, kwargs",

[

([], {}),

([LiteralDoc("arg")], {}),

([LiteralDoc("arg"), IdDoc("x")], {}),

([], {"x": LiteralDoc("x")}),

([], {"x": LiteralDoc("x"), "y": LiteralDoc("y")}),

([LiteralDoc("arg")], {"x": LiteralDoc("x"), "y": LiteralDoc("y")}),

([LiteralDoc("arg"), IdDoc("x")], {"x": LiteralDoc("x"), "y": LiteralDoc("y")}),

],

)

def test_expr_doc_call_with(args, kwargs):

target = IdDoc("x")

doc = target.call(*args, **kwargs)

assert doc.callee == target

assert list(doc.args) == args

assert dict(zip(doc.kwargs_keys, doc.kwargs_values)) == kwargs

@pytest.mark.parametrize(

"stmts",

[

[],

[ExprStmtDoc(IdDoc("x"))],

[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],

],

)

def test_stmt_block_doc(stmts):

doc = StmtBlockDoc(stmts)

assert list(doc.stmts) == stmts

@pytest.mark.parametrize(

"lhs, rhs, annotation",

[

(IdDoc("x"), IdDoc("y"), None),

(IdDoc("x"), None, IdDoc("int")),

(IdDoc("x"), IdDoc("y"), IdDoc("int")),

],

)

def test_assign_doc(lhs, rhs, annotation):

doc = AssignDoc(lhs, rhs, annotation)

assert doc.lhs == lhs

assert doc.rhs == rhs

assert doc.annotation == annotation

@pytest.mark.parametrize(

"lhs, rhs, annotation",

[

(IdDoc("x"), None, None),

(TupleDoc([IdDoc("x"), IdDoc("y")]), None, IdDoc("int")),

(TupleDoc([IdDoc("x"), IdDoc("y")]), IdDoc("u"), IdDoc("int")),

],

)

def test_invalid_assign_doc(lhs, rhs, annotation):

with pytest.raises(ValueError) as e:

AssignDoc(lhs, rhs, annotation)

@pytest.mark.parametrize(

"else_branch",

[

[],

[ExprStmtDoc(IdDoc("x"))],

[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],

],

)

@pytest.mark.parametrize(

"then_branch",

[

[],

[ExprStmtDoc(IdDoc("x"))],

[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],

],

)

def test_if_doc(then_branch, else_branch):

predicate = IdDoc("x")

if not then_branch and not else_branch:

with pytest.raises(ValueError) as e:

IfDoc(predicate, then_branch, else_branch)

return

else:

doc = IfDoc(predicate, then_branch, else_branch)

assert doc.predicate == predicate

assert list(doc.then_branch) == then_branch

assert list(doc.else_branch) == else_branch

@pytest.mark.parametrize(

"body",

[

[],

[ExprStmtDoc(IdDoc("x"))],

[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],

],

)

def test_while_doc(body):

predicate = IdDoc("x")

doc = WhileDoc(predicate, body)

assert doc.predicate == predicate

assert list(doc.body) == body

@pytest.mark.parametrize(

"body",

[

[],

[ExprStmtDoc(IdDoc("x"))],

[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],

],

)

def test_for_doc(body):

lhs = IdDoc("x")

rhs = IdDoc("y")

doc = ForDoc(lhs, rhs, body)

assert doc.lhs == lhs

assert doc.rhs == rhs

assert list(doc.body) == body

@pytest.mark.parametrize(

"lhs",

[

None,

IdDoc("x"),

],

)

@pytest.mark.parametrize(

"body",

[

[],

[ExprStmtDoc(IdDoc("x"))],

[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],

],

)

def test_scope_doc(lhs, body):

rhs = IdDoc("y")

doc = ScopeDoc(lhs, rhs, body)

assert doc.lhs == lhs

assert doc.rhs == rhs

assert list(doc.body) == body

def test_expr_stmt_doc():

expr = IdDoc("x")

doc = ExprStmtDoc(expr)

assert doc.expr == expr

@pytest.mark.parametrize(

"msg",

[

None,

LiteralDoc("msg"),

],

)

def test_assert_doc(msg):

test = IdDoc("x")

doc = AssertDoc(test, msg)

assert doc.test == test

assert doc.msg == msg

def test_return_doc():

value = IdDoc("x")

doc = ReturnDoc(value)

assert doc.value == value

@pytest.mark.parametrize(

"args",

[

[],

[AssignDoc(IdDoc("x"), None, IdDoc("int"))],

[

AssignDoc(IdDoc("x"), None, IdDoc("int")),

AssignDoc(IdDoc("y"), LiteralDoc(1), IdDoc("int")),

],

],

)

@pytest.mark.parametrize(

"decorators",

[

[],

[IdDoc("test")],

[IdDoc("test"), IdDoc("test2")],

],

)

@pytest.mark.parametrize(

"return_type",

[

None,

LiteralDoc(None),

],

)

@pytest.mark.parametrize(

"body",

[

[],

[ExprStmtDoc(IdDoc("x"))],

[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],

],

)

def test_function_doc(args, decorators, return_type, body):

name = IdDoc("name")

doc = FunctionDoc(name, args, decorators, return_type, body)

assert doc.name == name

assert list(doc.args) == args

assert list(doc.decorators) == decorators

assert doc.return_type == return_type

assert list(doc.body) == body

@pytest.mark.parametrize(

"decorators",

[

[],

[IdDoc("test")],

[IdDoc("test"), IdDoc("test2")],

],

)

@pytest.mark.parametrize(

"body",

[

[],

[ExprStmtDoc(IdDoc("x"))],

[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],

],

)

def test_class_doc(decorators, body):

name = IdDoc("name")

doc = ClassDoc(name, decorators, body)

assert doc.name == name

assert list(doc.decorators) == decorators

assert list(doc.body) == body

@pytest.mark.parametrize(

"comment",

[

"",

"test comment 1",

"test comment 1\ntest comment 1",

],

)

def test_comment_doc(comment):

doc = CommentDoc(comment)

assert doc.comment == comment

@pytest.mark.parametrize(

"comment",

[

"",

"test comment 1",

"test comment 1\ntest comment 1",

],

)

def test_doc_string_doc(comment):

doc = DocStringDoc(comment)

assert doc.comment == comment

def test_stmt_doc_comment():

doc = ExprStmtDoc(IdDoc("x"))

assert doc.comment is None

comment = "test comment"

doc.comment = comment

# Make sure the previous statement doesn't set attribute

# as if it's an ordinary Python object (__slots__ enforces this).

assert not hasattr(doc, "__dict__") or "comment" not in doc.__dict__

assert doc.comment == comment

def test_doc_source_paths():

doc = IdDoc("x")

assert len(doc.source_paths) == 0

source_paths = [AccessPath.root(), AccessPath.root().attr("x")]

doc.source_paths = source_paths

# This should triggers the __getattr__ and gets a tvm_ffi.Array

assert not isinstance(doc.source_paths, list)

assert list(doc.source_paths) == source_paths

doc.source_paths = []

assert len(doc.source_paths) == 0

if __name__ == "__main__":

tvm.testing.main()