diff --git a/pystencils/display_utils.py b/pystencils/display_utils.py index f6c32ac88ae68d684e860b33c0e5185ccc030e4e..e83067e90cba9585c1c31e88ac335e4205edb1ef 100644 --- a/pystencils/display_utils.py +++ b/pystencils/display_utils.py @@ -70,6 +70,32 @@ def get_code_obj(ast: Union[KernelFunction, KernelWrapper], custom_backend=None) return CodeDisplay(ast) +def get_assembly_obj(ast: Union[KernelFunction, KernelWrapper]): + """Returns an object to display generated assembly code (C/C++) + + Can either be displayed as HTML in Jupyter notebooks or printed as normal string. + """ + from pystencils.backends.cbackend import generate_c + + if isinstance(ast, KernelWrapper): + ast = ast.ast + code = generate_c(ast, dialect=ast.backend) + + class AssemlbyDisplay: + def __init__(self, assembly): + self.assembly = assembly + + def _repr_html(self): + return '' + + def __str__(self): + return '' + + def __repr__(self): + return '' + return AssemlbyDisplay(assembly) + + def get_code_str(ast, custom_backend=None): return str(get_code_obj(ast, custom_backend)) @@ -102,3 +128,20 @@ def show_code(ast: Union[KernelFunction, KernelWrapper], custom_backend=None): console.print(syntax) except ImportError: print(code) + + +def show_assembly(ast: Union[KernelFunction, KernelWrapper]): + code = get_assembly_obj(ast) + + if _isnotebook(): + from IPython.display import display + display(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/pystencils_tests/test_show_assembly.py b/pystencils_tests/test_show_assembly.py new file mode 100644 index 0000000000000000000000000000000000000000..74b2b06f15514f4b3bf5a08ae71fd5865e83867b --- /dev/null +++ b/pystencils_tests/test_show_assembly.py @@ -0,0 +1,8 @@ +import pytest +import sympy as sp + +import pystencils as ps + + +def test_simple(): + pass