pystencils/data_types.py

@@ -54,6 +54,7 @@ class cast_func(sp.Function):
         # rhs = cast_func(0, 'int')
         # print( sp.Ne(lhs, rhs) ) # would give true if all cast_funcs are booleans
         # -> thus a separate class boolean_cast_func is introduced
+        assert isinstance(args[1], Type), "Second argument of cast_func must be of type pymbolics.data_type.Type"
         if isinstance(args[0], Boolean):
             cls = boolean_cast_func
         return sp.Function.__new__(cls, *args, **kwargs)

pystencils_tests/test_address_of.py

@@ -3,19 +3,18 @@
 Test of pystencils.data_types.address_of
 """
 
-from pystencils.data_types import address_of, cast_func, PointerType
 import pystencils
+from pystencils.data_types import PointerType, address_of, cast_func, create_type
 from pystencils.simp.simplifications import sympy_cse
-import sympy
 
 
 def test_address_of():
     x, y = pystencils.fields('x,y: int64[2d]')
-    s = pystencils.TypedSymbol('s', PointerType('int64'))
+    s = pystencils.TypedSymbol('s', PointerType(create_type('int64')))
 
     assignments = pystencils.AssignmentCollection({
         s: address_of(x[0, 0]),
-        y[0, 0]: cast_func(s, 'int64')
+        y[0, 0]: cast_func(s, create_type('int64'))
     }, {})
 
     ast = pystencils.create_kernel(assignments)
@@ -23,7 +22,7 @@ def test_address_of():
     print(code)
 
     assignments = pystencils.AssignmentCollection({
-        y[0, 0]: cast_func(address_of(x[0, 0]), 'int64')
+        y[0, 0]: cast_func(address_of(x[0, 0]), create_type('int64'))
     }, {})
 
     ast = pystencils.create_kernel(assignments)
@@ -33,11 +32,11 @@ def test_address_of():
 
 def test_address_of_with_cse():
     x, y = pystencils.fields('x,y: int64[2d]')
-    s = pystencils.TypedSymbol('s', PointerType('int64'))
+    s = pystencils.TypedSymbol('s', PointerType(create_type('int64')))
 
     assignments = pystencils.AssignmentCollection({
-        y[0, 0]: cast_func(address_of(x[0, 0]), 'int64'),
-        x[0, 0]: cast_func(address_of(x[0, 0]), 'int64') + 1
+        y[0, 0]: cast_func(address_of(x[0, 0]), create_type('int64')) + s,
+        x[0, 0]: cast_func(address_of(x[0, 0]), create_type('int64')) + 1
     }, {})
 
     ast = pystencils.create_kernel(assignments)