Class Parser

python/tvm/script/parser/core/parser.py:352–952

The TVMScript parser Parameters ---------- diag : Diagnostics The diagnostics for error reporting. dispatch_tokens : List[str] The list of dispatching tokens to dispatching parsing method of different IRs and different doc AST structure. var_table : Var

class Parser(doc.NodeVisitor):
    """The TVMScript parser

    Parameters
    ----------
    diag : Diagnostics
        The diagnostics for error reporting.

    dispatch_tokens : List[str]
        The list of dispatching tokens to dispatching parsing method
        of different IRs and different doc AST structure.

    var_table : VarTable
        The variable table for parsing.
    """

    diag: Diagnostics
    dispatch_tokens: list[str]
    function_annotations: dict[str, dict[str, Any]] | None
    var_table: VarTable
    inside_function: bool  # whether we are within a function
    current_class: str | None = None  # current class being parsed
    base_py_module_context: bool = False  # whether current class inherits from BasePyModule

    def __init__(
        self,
        source: Source,
        function_annotations: dict[str, dict[str, Any]],
    ) -> None:
        self.diag = Diagnostics(source)
        self.dispatch_tokens = ["default"]
        self.function_annotations = function_annotations
        self.var_table = VarTable()
        self.inside_function = False

    def parse(self, extra_vars: dict[str, Any] | None = None) -> Any:
        """The main parse method for parser.

        Parameters
        ----------
        extra_vars : Optional[Dict[str, Any]]
            The optional global value table for parsing.

        Returns
        -------
        res : Any
            The doc AST node visiting result.
        """
        if extra_vars is None:
            extra_vars = {}
        with self.var_table.with_frame():
            for k, v in extra_vars.items():
                self.var_table.add(k, v)
            node = self.diag.source.as_ast()
            self.visit(node)

    def get_dispatch_token(self, node: doc.FunctionDef) -> str:
        if not isinstance(node, doc.FunctionDef):