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 <jaway....@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)))