Function visit_with

python/tvm/tirx/script/parser/parser.py:569–594

The with visiting method for tirx. Parameters ---------- self : Parser The visiting parser. node : doc.With The doc AST with node.

(self: Parser, node: doc.With)

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

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

    node : doc.With
        The doc AST with node.
    """
    with contextlib.ExitStack() as stack:
        stack.enter_context(self.var_table.with_frame())
        for item in node.items:
            frame = self.eval_expr(item.context_expr)
            if not isinstance(frame, Frame) and not (
                hasattr(frame, "__enter__") and hasattr(frame, "__exit__")
            ):
                self.report_error(
                    item.context_expr,
                    "Invalid context expression in the with-statement.",
                )
            rhs = stack.enter_context(frame)
            if item.optional_vars is not None:
                self.eval_assign(target=item.optional_vars, source=rhs, bind_value=bind_with_value)
        self.visit_body(node.body)


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