WIP

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)

pystencils_tests/test_show_assembly.py

+import pytest
+import sympy as sp
+
+import pystencils as ps
+
+
+def test_simple():
+    pass