[tvm] branch main updated: [TVMScript] More accurate hints for ImportError (#13662)

[masahi]

This is an automated email from the ASF dual-hosted git repository.

masahi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new dc3fc36d0b [TVMScript] More accurate hints for ImportError (#13662)
dc3fc36d0b is described below

commit dc3fc36d0bd10d0c261ee6b5c99ddcd631d7f740
Author: Hongyu Cai <h.t...@hotmail.com>
AuthorDate: Wed Dec 28 06:10:50 2022 -0500

    [TVMScript] More accurate hints for ImportError (#13662)
---
 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 d12f6c2767..5cf28fff3a 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,
             )