Fix #6: let show_code display/print code

- `get_code` will keep the previous behavior of `show_code`
- `get_code_str` will be equivalent to previous `str(show_code(...))`

pystencils/__init__.py

@@ -5,7 +5,7 @@ from . import stencil as stencil
 from .assignment import Assignment, assignment_from_stencil
 from .data_types import TypedSymbol
 from .datahandling import create_data_handling
-from .display_utils import show_code, to_dot
+from .display_utils import show_code, get_code, get_code_str, to_dot
 from .field import Field, FieldType, fields
 from .kernel_decorator import kernel
 from .kernelcreation import create_indexed_kernel, create_kernel, create_staggered_kernel
@@ -25,7 +25,7 @@ __all__ = ['Field', 'FieldType', 'fields',
            'TypedSymbol',
            'make_slice',
            'create_kernel', 'create_indexed_kernel', 'create_staggered_kernel',
-           'show_code', 'to_dot',
+           'show_code', 'to_dot', 'get_code', 'get_code_str',
            'AssignmentCollection',
            'Assignment',
            'assignment_from_stencil',

pystencils/display_utils.py

-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Union
 
 import sympy as sp
 
@@ -35,10 +35,10 @@ def highlight_cpp(code: str):
     return HTML(highlight(code, CppLexer(), HtmlFormatter()))
 
 
-def show_code(ast: KernelFunction, custom_backend=None):
+def get_code(ast: Union[KernelFunction, KernelWrapper], custom_backend=None):
     """Returns an object to display generated code (C/C++ or CUDA)
 
-    Can either  be displayed as HTML in Jupyter notebooks or printed as normal string.
+    Can either be displayed as HTML in Jupyter notebooks or printed as normal string.
     """
     from pystencils.backends.cbackend import generate_c
 
@@ -65,3 +65,17 @@ def show_code(ast: KernelFunction, custom_backend=None):
         def __repr__(self):
             return generate_c(self.ast, dialect=dialect, custom_backend=custom_backend)
     return CodeDisplay(ast)
+
+
+def get_code_str(ast, custom_backend=None):
+    return str(get_code(ast, custom_backend))
+
+
+def show_code(ast: Union[KernelFunction, KernelWrapper], custom_backend=None):
+    code = get_code(ast, custom_backend)
+
+    try:
+        from IPython.display import display
+        display(code)
+    except Exception:
+        print(code)

pystencils/kernel_wrapper.py

@@ -16,4 +16,4 @@ class KernelWrapper:
 
     @property
     def code(self):
-        return str(pystencils.show_code(self.ast))
+        return pystencils.get_code_str(self.ast)

pystencils/test_type_interference.py

@@ -18,8 +18,7 @@ def test_type_interference():
 
     ast = pystencils.create_kernel(assignments)
 
-    code = str(pystencils.show_code(ast))
-    print(code)
+    code = str(pystencils.get_code_str(ast))
     assert 'double a' in code
     assert 'uint16_t b' in code
     assert 'uint16_t f' in code

pystencils_tests/test_address_of.py

@@ -17,16 +17,14 @@ def test_address_of():
     }, {})
 
     ast = pystencils.create_kernel(assignments)
-    code = pystencils.show_code(ast)
-    print(code)
+    pystencils.show_code(ast)
 
     assignments = pystencils.AssignmentCollection({
         y[0, 0]: cast_func(address_of(x[0, 0]), create_type('int64'))
     }, {})
 
     ast = pystencils.create_kernel(assignments)
-    code = pystencils.show_code(ast)
-    print(code)
+    pystencils.show_code(ast)
 
 
 def test_address_of_with_cse():
@@ -39,9 +37,8 @@ def test_address_of_with_cse():
     }, {})
 
     ast = pystencils.create_kernel(assignments)
-    code = pystencils.show_code(ast)
+    pystencils.show_code(ast)
     assignments_cse = sympy_cse(assignments)
 
     ast = pystencils.create_kernel(assignments_cse)
-    code = pystencils.show_code(ast)
-    print(code)
+    pystencils.show_code(ast)

pystencils_tests/test_complex_numbers.py

@@ -52,7 +52,7 @@ def test_complex_numbers(assignment, scalar_dtypes, target):
     ast = pystencils.create_kernel(assignment,
                                    target=target,
                                    data_type=scalar_dtypes)
-    code = str(pystencils.show_code(ast))
+    code = pystencils.get_code_str(ast)
 
     print(code)
     assert "Not supported" not in code
@@ -95,7 +95,7 @@ def test_complex_numbers_64(assignment, target):
     ast = pystencils.create_kernel(assignment,
                                    target=target,
                                    data_type='double')
-    code = str(pystencils.show_code(ast))
+    code = pystencils.get_code_str(ast)
 
     print(code)
     assert "Not supported" not in code

pystencils_tests/test_conditional_field_access.py

@@ -63,7 +63,7 @@ def test_boundary_check(with_cse):
 
     print(assignments)
     kernel_checked = ps.create_kernel(assignments, ghost_layers=0).compile()
-    print(ps.show_code(kernel_checked))
+    ps.show_code(kernel_checked)
 
     # No SEGFAULT, please!!
     kernel_checked(f=f_arr, g=g_arr)

pystencils_tests/test_cuda_known_functions.py

@@ -18,7 +18,7 @@ def test_cuda_known_functions():
     })
 
     ast = pystencils.create_kernel(assignments, 'gpu')
-    print(pystencils.show_code(ast))
+    pystencils.show_code(ast)
     kernel = ast.compile()
     assert(kernel is not None)
 
@@ -32,7 +32,7 @@ def test_cuda_but_not_c():
     })
 
     ast = pystencils.create_kernel(assignments, 'cpu')
-    print(pystencils.show_code(ast))
+    pystencils.show_code(ast)
 
 
 def test_cuda_unknown():
@@ -43,5 +43,4 @@ def test_cuda_unknown():
     })
 
     ast = pystencils.create_kernel(assignments, 'gpu')
-    code = str(pystencils.show_code(ast))
-    print(code)
+    pystencils.show_code(ast)

pystencils_tests/test_custom_backends.py

 from subprocess import CalledProcessError
 
-import pycuda.driver
 import pytest
 import sympy
 
+import pycuda.driver
 import pystencils
 import pystencils.cpu.cpujit
 import pystencils.gpucuda.cudajit
@@ -32,11 +32,11 @@ def test_custom_backends():
         z[0, 0], x[0, 0] * sympy.log(x[0, 0] * y[0, 0]))], [])
 
     ast = pystencils.create_kernel(normal_assignments, target='cpu')
-    print(pystencils.show_code(ast, ScreamingBackend()))
+    pystencils.show_code(ast, ScreamingBackend())
     with pytest.raises(CalledProcessError):
         pystencils.cpu.cpujit.make_python_function(ast, custom_backend=ScreamingBackend())
 
     ast = pystencils.create_kernel(normal_assignments, target='gpu')
-    print(pystencils.show_code(ast, ScreamingGpuBackend()))
+    pystencils.show_code(ast, ScreamingGpuBackend())
     with pytest.raises(pycuda.driver.CompileError):
         pystencils.gpucuda.cudajit.make_python_function(ast, custom_backend=ScreamingGpuBackend())

pystencils_tests/test_fast_approximation.py

@@ -12,7 +12,7 @@ def test_fast_sqrt():
     assert len(insert_fast_sqrts(expr).atoms(fast_sqrt)) == 1
     assert len(insert_fast_sqrts([expr])[0].atoms(fast_sqrt)) == 1
     ast = ps.create_kernel(ps.Assignment(g[0, 0], insert_fast_sqrts(expr)), target='gpu')
-    code_str = str(ps.show_code(ast))
+    code_str = ps.get_code_str(ast)
     assert '__fsqrt_rn' in code_str
 
     expr = ps.Assignment(sp.Symbol("tmp"), 3 / sp.sqrt(f[0, 0] + f[1, 0]))
@@ -21,7 +21,7 @@ def test_fast_sqrt():
     ac = ps.AssignmentCollection([expr], [])
     assert len(insert_fast_sqrts(ac).main_assignments[0].atoms(fast_inv_sqrt)) == 1
     ast = ps.create_kernel(insert_fast_sqrts(ac), target='gpu')
-    code_str = str(ps.show_code(ast))
+    code_str = ps.get_code_str(ast)
     assert '__frsqrt_rn' in code_str
 
 
@@ -34,5 +34,5 @@ def test_fast_divisions():
     assert len(insert_fast_divisions(expr).atoms(fast_division)) == 1
 
     ast = ps.create_kernel(ps.Assignment(g[0, 0], insert_fast_divisions(expr)), target='gpu')
-    code_str = str(ps.show_code(ast))
+    code_str = ps.get_code_str(ast)
     assert '__fdividef' in code_str

pystencils_tests/test_interpolation.py

@@ -11,11 +11,11 @@ import itertools
 from os.path import dirname, join
 
 import numpy as np
-import pycuda.autoinit  # NOQA
-import pycuda.gpuarray as gpuarray
 import pytest
 import sympy
 
+import pycuda.autoinit  # NOQA
+import pycuda.gpuarray as gpuarray
 import pystencils
 from pystencils.interpolation_astnodes import LinearInterpolator
 from pystencils.spatial_coordinates import x_, y_
@@ -51,7 +51,7 @@ def test_interpolation():
     print(assignments)
     ast = pystencils.create_kernel(assignments)
     print(ast)
-    print(pystencils.show_code(ast))
+    pystencils.show_code(ast)
     kernel = ast.compile()
 
     pyconrad.imshow(lenna)
@@ -71,7 +71,7 @@ def test_scale_interpolation():
         print(assignments)
         ast = pystencils.create_kernel(assignments)
         print(ast)
-        print(pystencils.show_code(ast))
+        pystencils.show_code(ast)
         kernel = ast.compile()
 
         out = np.zeros_like(lenna)
@@ -102,7 +102,7 @@ def test_rotate_interpolation(address_mode):
     print(assignments)
     ast = pystencils.create_kernel(assignments)
     print(ast)
-    print(pystencils.show_code(ast))
+    pystencils.show_code(ast)
     kernel = ast.compile()
 
     out = np.zeros_like(lenna)
@@ -135,7 +135,7 @@ def test_rotate_interpolation_gpu(address_mode):
             print(assignments)
             ast = pystencils.create_kernel(assignments, target='gpu', use_textures_for_interpolation=use_textures)
             print(ast)
-            print(pystencils.show_code(ast))
+            pystencils.show_code(ast)
             kernel = ast.compile()
 
             out = gpuarray.zeros_like(lenna_gpu)
@@ -182,7 +182,7 @@ def test_shift_interpolation_gpu(address_mode):
             # print(assignments)
             ast = pystencils.create_kernel(assignments, target='gpu', use_textures_for_interpolation=use_textures)
             # print(ast)
-            print(pystencils.show_code(ast))
+            pystencils.show_code(ast)
             kernel = ast.compile()
 
             out = gpuarray.zeros_like(lenna_gpu)
@@ -209,7 +209,7 @@ def test_rotate_interpolation_size_change(address_mode):
     print(assignments)
     ast = pystencils.create_kernel(assignments)
     print(ast)
-    print(pystencils.show_code(ast))
+    pystencils.show_code(ast)
     kernel = ast.compile()
 
     out = np.zeros((100, 150), np.float64)
@@ -228,7 +228,7 @@ def test_field_interpolated(address_mode, target):
     print(assignments)
     ast = pystencils.create_kernel(assignments)
     print(ast)
-    print(pystencils.show_code(ast))
+    pystencils.show_code(ast)
     kernel = ast.compile()
 
     out = np.zeros_like(lenna)

pystencils_tests/test_jacobi_cbackend.py

 import numpy as np
 
-from pystencils import show_code
+from pystencils import get_code
 from pystencils.astnodes import Block, KernelFunction, SympyAssignment
 from pystencils.cpu import make_python_function
 from pystencils.field import Field
@@ -36,7 +36,7 @@ def test_jacobi_fixed_field_size():
     error = np.sum(np.abs(dst_field_py - dst_field_c))
     np.testing.assert_allclose(error, 0.0, atol=1e-13)
 
-    code_display = show_code(ast_node)
+    code_display = get_code(ast_node)
     assert 'for' in str(code_display)
     assert 'for' in code_display._repr_html_()
 

pystencils_tests/test_jacobi_llvm.py

@@ -52,7 +52,7 @@ def test_jacobi_fixed_field_size_gpu():
 
     jacobi = Assignment(d[0, 0], (f[1, 0] + f[-1, 0] + f[0, 1] + f[0, -1]) / 4)
     ast = create_kernel([jacobi], target='gpu')
-    print(show_code(ast))
+    show_code(ast)
 
     for x in range(1, size[0] - 1):
         for y in range(1, size[1] - 1):

pystencils_tests/test_opencl.py

@@ -27,10 +27,9 @@ def test_print_opencl():
 
     print(ast)
 
-    code = pystencils.show_code(ast, custom_backend=CudaBackend())
-    print(code)
+    pystencils.show_code(ast, custom_backend=CudaBackend())
 
-    opencl_code = pystencils.show_code(ast, custom_backend=OpenClBackend())
+    opencl_code = pystencils.get_code_str(ast, custom_backend=OpenClBackend())
     print(opencl_code)
 
     assert "__global double * RESTRICT const _data_x" in str(opencl_code)
@@ -108,10 +107,9 @@ def test_opencl_jit():
 
     print(ast)
 
-    code = pystencils.show_code(ast, custom_backend=CudaBackend())
-    print(code)
-    opencl_code = pystencils.show_code(ast, custom_backend=OpenClBackend())
-    print(opencl_code)
+    pystencils.show_code(ast, custom_backend=CudaBackend())
+
+    pystencils.show_code(ast, custom_backend=OpenClBackend())
 
     cuda_kernel = ast.compile()
     assert cuda_kernel is not None

pystencils_tests/test_source_code_comment.py

@@ -27,4 +27,5 @@ def test_source_code_comment():
     print(ast)
     compiled = ast.compile()
     assert compiled is not None
-    print(pystencils.show_code(ast))
+
+    pystencils.show_code(ast)

pystencils_tests/test_sum_prod.py

@@ -120,11 +120,9 @@ def test_prod_var_limit():
     })
 
     ast = pystencils.create_kernel(assignments)
-    code = str(pystencils.show_code(ast))
+    pystencils.show_code(ast)
     kernel = ast.compile()
 
-    print(code)
-
     array = np.zeros((10,), np.int64)
 
     kernel(x=array, limit=100)