diff --git a/pystencils/display_utils.py b/pystencils/display_utils.py index 0bdad7ac16ba6c600cd66e5c82721c8ac46bb0dd..cb09197bf175b43c4896072693ea11253976691f 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 83a76f870f02f64d8a371f294c1224c15c919061..8eeeff7b2e4a324feb7dc31fc9b8605254ae2da0 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'],