From 8581278b81ef5bf889f4aa80098934c7afee1df3 Mon Sep 17 00:00:00 2001 From: Junru Shao Date: Fri, 1 Jul 2022 18:29:33 -0700 Subject: [PATCH] add rendering of error in dispatch (#58) --- python/tvm/script/parse/parser.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/python/tvm/script/parse/parser.py b/python/tvm/script/parse/parser.py index 94bb303f5c77..91c769764a23 100644 --- a/python/tvm/script/parse/parser.py +++ b/python/tvm/script/parse/parser.py @@ -33,12 +33,23 @@ } +def _dispatch_wrapper(func: dispatch.ParseMethod) -> dispatch.ParseMethod: + def _wrapper(self: Parser, node: doc.AST) -> None: + try: + return func(self, node) + except Exception as e: + self.report_error(node, str(e)) + raise + + return _wrapper + + def _dispatch(self: "Parser", type_name: str) -> dispatch.ParseMethod: for token in [self.dispatch_tokens[-1], "default"]: func = dispatch.get(token=token, type_name=type_name, default=None) if func is not None: - return func - return lambda self, node: self.generic_visit(node) + return _dispatch_wrapper(func) + return _dispatch_wrapper(lambda self, node: self.generic_visit(node)) class Parser(doc.NodeVisitor):