Use type annotations in generated code

psydac/api/ast/parser.py

@@ -76,7 +76,8 @@
 from psydac.api.ast.utilities import variables, math_atoms_as_str, get_name
 from psydac.api.utilities     import flatten
 from psydac.api.ast.utilities import build_pythran_types_header
-from psydac.api.ast.utilities import build_pyccel_types_decorator
+#from psydac.api.ast.utilities import build_pyccel_types_decorator
+from psydac.api.ast.utilities import build_pyccel_type_annotations
 
 #==============================================================================
 # TODO move it
@@ -560,10 +561,16 @@ def _visit_DefNode(self, expr, **kwargs):
         results       = [self._visit(a) for a in expr.results]
 
         if self.backend['name'] == 'pyccel':
-            a = [String(str(i)) for i in build_pyccel_types_decorator(arguments)]
-            decorators = {'types': Function('types')(*a)}
+#            a = [String(str(i)) for i in build_pyccel_types_decorator(arguments)]
+#            decorators = {'types': Function('types')(*a)}
+            arguments = build_pyccel_type_annotations(arguments)
+#            print('-'*30)
+#            for a in arguments:
+#                print(self._visit(a))
+#            print('-'*30)
+            decorators = {}
         elif self.backend['name'] == 'pythran':
-            header = build_pythran_types_header(name, arguments)
+            header = build_pythran_types_header(name, arguments) # Could this work??
         else:
             decorators = {}
 

psydac/api/ast/utilities.py

@@ -31,6 +31,7 @@
 from psydac.pyccel.ast.core import _atomic
 from psydac.pyccel.ast.core import Comment
 from psydac.pyccel.ast.core import String
+from psydac.pyccel.ast.core import AnnotatedArgument
 
 #==============================================================================
 def random_string( n ):
@@ -815,6 +816,47 @@ def literal(s):
     else:
         raise TypeError('Expecting a string')
 
+#==============================================================================
+def build_pyccel_type_annotations(args, order=None):
+
+    new_args = []
+
+    for a in args:
+        if isinstance(a, Variable):
+            rank  = a.rank
+            dtype = a.dtype.name.lower()
+
+        elif isinstance(a, IndexedVariable):
+            rank  = a.rank
+            dtype = a.dtype.name.lower()
+
+        elif isinstance(a, Constant):
+            rank = 0
+            if a.is_integer:
+                dtype = 'int'
+            elif a.is_real:
+                dtype = 'float'
+            elif a.is_complex:
+                dtype = 'complex'
+            else:
+                raise TypeError(f"The Constant {a} don't have any information about the type of the variable.\n"
+                                f"Please create the Constant like this Constant('{a}', real=True), Constant('{a}', complex=True) or Constant('{a}', integer=True).")
+
+        else:
+            raise TypeError('unexpected type for {}'.format(a))
+
+        if rank > 0:
+            shape = ','.join(':' * rank)
+            dtype = '{dtype}[{shape}]'.format(dtype=dtype, shape=shape)
+            if order and rank > 1:
+                dtype = "{dtype}(order={ordering})".format(dtype=dtype, ordering=order)
+
+        dtype = String(dtype)
+        new_a = AnnotatedArgument(a, dtype)
+        new_args.append(new_a)
+
+    return new_args
+
 #==============================================================================
 def build_pyccel_types_decorator(args, order=None):
     """

psydac/pyccel/ast/core.py

@@ -5705,3 +5705,18 @@ def fprint(self, printer, lhs = None):
             code_init = 'size({0}, {1})'.format(init_value, index)
 
         return code_init
+
+
+class AnnotatedArgument(Symbol, PyccelAstNode):
+
+    def __new__(cls, name, annotation):
+        return Basic.__new__(cls, name, annotation)
+
+    @property
+    def name(self):
+        return self.args[0]
+
+    @property
+    def annotation(self):
+        return self.args[1]
+

psydac/pyccel/codegen/printing/pycode.py

@@ -363,6 +363,11 @@ def _print_PyccelNot(self, expr):
         a = self._print(expr.args[0])
         return 'not {}'.format(a)
 
+    def _print_AnnotatedArgument(self, expr):
+        name       = self._print(expr.name)
+        annotation = self._print(expr.annotation)
+        return f'{name} : {annotation}'
+
 #==============================================================================
 def pycode(expr, **settings):
     """ Converts an expr to a string of Python code