From 26445cdf1e71658f29d41328989fe980b45c721a Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Thu, 30 Jan 2020 10:11:30 +0100
Subject: [PATCH] Use `rich` for syntax highlighting of `show_code` also in
terminal
---
pystencils/display_utils.py | 26 +++++++++++++++++++++++---
setup.py | 2 +-
2 files changed, 24 insertions(+), 4 deletions(-)
diff --git a/pystencils/display_utils.py b/pystencils/display_utils.py
index 0bdad7ac1..cb09197bf 100644
--- a/pystencils/display_utils.py
+++ b/pystencils/display_utils.py
@@ -71,11 +71,31 @@ def get_code_str(ast, custom_backend=None):
return str(get_code_obj(ast, custom_backend))
+def _isnotebook():
+ try:
+ shell = get_ipython().__class__.__name__
+ if shell == 'ZMQInteractiveShell':
+ return True # Jupyter notebook or qtconsole
+ elif shell == 'TerminalInteractiveShell':
+ return False # Terminal running IPython
+ else:
+ return False # Other type (?)
+ except NameError:
+ return False
+
+
def show_code(ast: Union[KernelFunction, KernelWrapper], custom_backend=None):
code = get_code_obj(ast, custom_backend)
- try:
+ if _isnotebook():
from IPython.display import display
display(code)
- except Exception:
- print(code)
+ else:
+ try:
+ import rich.syntax
+ import rich.console
+ syntax = rich.syntax.Syntax(str(code), "c++", theme="monokai", line_numbers=True)
+ console = rich.console.Console()
+ console.print(syntax)
+ except ImportError:
+ print(code)
diff --git a/setup.py b/setup.py
index 83a76f870..8eeeff7b2 100644
--- a/setup.py
+++ b/setup.py
@@ -112,7 +112,7 @@ setup(name='pystencils',
'opencl': ['pyopencl'],
'alltrafos': ['islpy', 'py-cpuinfo'],
'bench_db': ['blitzdb', 'pymongo', 'pandas'],
- 'interactive': ['matplotlib', 'ipy_table', 'imageio', 'jupyter', 'pyevtk'],
+ 'interactive': ['matplotlib', 'ipy_table', 'imageio', 'jupyter', 'pyevtk', 'rich'],
'autodiff': ['pystencils-autodiff'],
'doc': ['sphinx', 'sphinx_rtd_theme', 'nbsphinx',
'sphinxcontrib-bibtex', 'sphinx_autodoc_typehints', 'pandoc'],
--
GitLab