From c1bb44a74c3e3f7bbb2c17e113a33327e460851c Mon Sep 17 00:00:00 2001 From: Hongyu Cai Date: Mon, 26 Dec 2022 11:07:38 -0500 Subject: [PATCH] [TVMScript] More accurate hints for ImportError --- python/tvm/script/highlight.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/python/tvm/script/highlight.py b/python/tvm/script/highlight.py index d12f6c276767..5cf28fff3a4b 100644 --- a/python/tvm/script/highlight.py +++ b/python/tvm/script/highlight.py @@ -160,13 +160,20 @@ def _get_pygments_style( if version.parse(pygments.__version__) < version.parse("2.4.0"): raise ImportError("Required Pygments version >= 2.4.0 but got " + pygments.__version__) except ImportError as err: + if err.name == "packaging": + name = "packaging" + elif err.name == "pygments": + name = "Pygments>=2.4.0" + else: + raise ValueError(f'Package "{err.name}" should not be used') + with warnings.catch_warnings(): warnings.simplefilter("once", UserWarning) - install_cmd = sys.executable + ' -m pip install "Pygments>=2.4.0" --upgrade --user' + install_cmd = sys.executable + f' -m pip install "{name}" --upgrade --user' warnings.warn( str(err) + "\n" - + "To print highlighted TVM script, please install Pygments:\n" + + f"To print highlighted TVM script, please install {name}:\n" + install_cmd, category=UserWarning, )