Function visit_assert

python/tvm/tirx/script/parser/parser.py:804–849

The assert visiting method for tirx. Parameters ---------- self : Parser The visiting parser. node : doc.Assert The doc AST assert node. The assert message can be either: - A plain string: ``assert cond, "message"`` - A tuple of (kind, [parts...]): ``as

(self: Parser, node: doc.Assert)

@dispatch.register(token="tirx", type_name="Assert")
def visit_assert(self: Parser, node: doc.Assert) -> None:
    """The assert visiting method for tirx.

    Parameters
    ----------
    self : Parser
        The visiting parser.

    node : doc.Assert
        The doc AST assert node.

    The assert message can be either:
    - A plain string: ``assert cond, "message"``
    - A tuple of (kind, [parts...]): ``assert cond, ("ValueError", ["part0", "part1"])``
    """
    cond = self.eval_expr(node.test)
    msg = self.eval_expr(node.msg)

    kind = "RuntimeError"
    message = msg

    if isinstance(msg, tuple):
        if len(msg) != 2:
            self.report_error(
                node,
                f"Assert message tuple must have exactly 2 elements (kind, [parts...]), "
                f"got {len(msg)} elements",
            )
        kind_str, parts = msg
        if isinstance(kind_str, tvm.tirx.StringImm):
            kind_str = kind_str.value
        if not isinstance(kind_str, str):
            self.report_error(
                node,
                f"Assert message tuple first element must be a string (error kind like "
                f'"ValueError"), got {type(kind_str).__name__}',
            )
        kind = kind_str
        message = parts

    if isinstance(message, list | tuple):
        message = [p.value if isinstance(p, tvm.tirx.StringImm) else str(p) for p in message]

    frame = T.Assert(cond, message, error_kind=kind)
    frame.add_callback(partial(frame.__exit__, None, None, None))
    frame.__enter__()


@dispatch.register(token="tirx", type_name="Return")