[tvm] branch main updated: [UX][TVMSciprt] Use HTML formatter in notebook environments (#12240)

[ju...@apache.org]

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

junrushao 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 e756980b41 [UX][TVMSciprt] Use HTML formatter in notebook environments (#12240)
e756980b41 is described below

commit e756980b41c339ad26ea50315eb29073504b3796
Author: Jiawei Liu <ja...@gmail.com>
AuthorDate: Sat Jul 30 16:33:45 2022 -0500

    [UX][TVMSciprt] Use HTML formatter in notebook environments (#12240)
    
    Previously we use ANSI color sequences to highlight TVM script. In jupyter notebook environments, such color sequence will be recoginized and translated to corresponding HTML to display things.
    
    This works fine for most notebook environments (including the jupyter notebook and the VS Code plugin). Recently, thanks to @tqchen, we found that Google Colab does not well support ansi color sequence for 24-bit colors (`JupyterLight` and `VSCDark`) that all its displayed colors are unexpectedly black/gray/white. To also bring highlighting in Colab, in this PR, we directly render the highlighted code with HTML when a notebook environment is detected.
---
 python/tvm/script/highlight.py | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/python/tvm/script/highlight.py b/python/tvm/script/highlight.py
index 03476ba60c..5a9c69a0ff 100644
--- a/python/tvm/script/highlight.py
+++ b/python/tvm/script/highlight.py
@@ -50,7 +50,7 @@ def cprint(printable: Union[IRModule, PrimFunc], style: Optional[str] = None) ->
         import pygments
         from pygments import highlight
         from pygments.lexers.python import Python3Lexer
-        from pygments.formatters import Terminal256Formatter
+        from pygments.formatters import Terminal256Formatter, HtmlFormatter
         from pygments.style import Style
         from pygments.token import Keyword, Name, Comment, String, Number, Operator
         from packaging import version
@@ -72,8 +72,9 @@ def cprint(printable: Union[IRModule, PrimFunc], style: Optional[str] = None) ->
     else:
 
         class JupyterLight(Style):
-            """A Jupyter-Notebook-like Pygments style configuration (aka. "dark")"""
+            """A Jupyter-Notebook-like Pygments style configuration (aka. "light")"""
 
+            background_color = ""
             styles = {
                 Keyword: "bold #008000",
                 Keyword.Type: "nobold #008000",
@@ -90,6 +91,7 @@ def cprint(printable: Union[IRModule, PrimFunc], style: Optional[str] = None) ->
         class VSCDark(Style):
             """A VSCode-Dark-like Pygments style configuration (aka. "dark")"""
 
+            background_color = ""
             styles = {
                 Keyword: "bold #c586c0",
                 Keyword.Type: "#82aaff",
@@ -107,6 +109,7 @@ def cprint(printable: Union[IRModule, PrimFunc], style: Optional[str] = None) ->
         class AnsiTerminalDefault(Style):
             """The default style for terminal display with ANSI colors (aka. "ansi")"""
 
+            background_color = ""
             styles = {
                 Keyword: "bold ansigreen",
                 Keyword.Type: "nobold ansigreen",
@@ -120,12 +123,11 @@ def cprint(printable: Union[IRModule, PrimFunc], style: Optional[str] = None) ->
                 Comment: "italic ansibrightblack",
             }
 
+        is_in_notebook = "ipykernel" in sys.modules  # in notebook env (support html display).
+
         if style is None:
             # choose style automatically according to the environment:
-            if "ipykernel" in sys.modules:  # in notebook env.
-                style = JupyterLight
-            else:  # in a terminal or something.
-                style = AnsiTerminalDefault
+            style = JupyterLight if is_in_notebook else AnsiTerminalDefault
         elif style == "light":
             style = JupyterLight
         elif style == "dark":
@@ -133,4 +135,12 @@ def cprint(printable: Union[IRModule, PrimFunc], style: Optional[str] = None) ->
         elif style == "ansi":
             style = AnsiTerminalDefault
 
-        print(highlight(printable.script(), Python3Lexer(), Terminal256Formatter(style=style)))
+        if is_in_notebook:  # print with HTML display
+            from IPython.display import display, HTML  # pylint: disable=import-outside-toplevel
+
+            formatter = HtmlFormatter(style=JupyterLight)
+            formatter.noclasses = True  # inline styles
+            html = highlight(printable.script(), Python3Lexer(), formatter)
+            display(HTML(html))
+        else:
+            print(highlight(printable.script(), Python3Lexer(), Terminal256Formatter(style=style)))