Add Fortran support for dict initialisation (pyccel#2005)

Add Fortran support for dict initialisation in a similar way to lists.
Fixes pyccel#1944

Author: EmilyBourne

pyccel/ast/builtins.py

@@ -908,6 +908,9 @@ def __repr__(self):
         args = ', '.join(f'{repr(k)}: {repr(v)}' for k,v in self)
         return f'PythonDict({args})'
 
+    def __len__(self):
+        return len(self._keys)
+
     @property
     def keys(self):
         """

pyccel/ast/low_level_tools.py

@@ -11,6 +11,7 @@
 from .datatypes import PyccelType
 
 __all__ = ('IteratorType',
+           'PairType',
            'MacroDefinition')
 
 #------------------------------------------------------------------------------
@@ -40,6 +41,51 @@ def iterable_type(self):
         """
         return self._iterable_type
 
+#------------------------------------------------------------------------------
+class PairType(PyccelType, metaclass=ArgumentSingleton):
+    """
+    The type of an element of a dictionary type.
+
+    The type of an element of a dictionary type.
+
+    Parameters
+    ----------
+    key_type : PyccelType
+        The type of the keys of the homogeneous dictionary.
+    value_type : PyccelType
+        The type of the values of the homogeneous dictionary.
+    """
+    __slots__ = ('_key_type', '_value_type')
+    _name = 'pair'
+    _container_rank = 0
+    _order = None
+
+    def __init__(self, key_type, value_type):
+        self._key_type = key_type
+        self._value_type = value_type
+        super().__init__()
+
+    @property
+    def key_type(self):
+        """
+        The type of the keys of the object.
+
+        The type of the keys of the object.
+        """
+        return self._key_type
+
+    @property
+    def value_type(self):
+        """
+        The type of the values of the object.
+
+        The type of the values of the object.
+        """
+        return self._value_type
+
+    def __str__(self):
+        return f'pair[{self._key_type}, {self._value_type}]'
+
 #------------------------------------------------------------------------------
 class MacroDefinition(PyccelAstNode):
     """

pyccel/codegen/printing/fcode.py

@@ -51,7 +51,7 @@
 from pyccel.ast.literals  import LiteralTrue, LiteralFalse, LiteralString
 from pyccel.ast.literals  import Nil
 
-from pyccel.ast.low_level_tools  import MacroDefinition, IteratorType
+from pyccel.ast.low_level_tools  import MacroDefinition, IteratorType, PairType
 
 from pyccel.ast.mathext  import math_constants
 
@@ -584,6 +584,37 @@ def _build_gFTL_module(self, expr_type):
                     raise NotImplementedError("Support for sets of types which define their own < operator is not yet implemented")
                 imports_and_macros.append(MacroDefinition('Set', expr_type))
                 imports_and_macros.append(MacroDefinition('SetIterator', IteratorType(expr_type)))
+            elif isinstance(expr_type, DictType):
+                include = Import(LiteralString('map/template.inc'), Module('_', (), ()))
+                key_type = expr_type.key_type
+                value_type = expr_type.value_type
+                imports_and_macros = []
+                if isinstance(key_type, FixedSizeNumericType):
+                    tmpVar_x = Variable(key_type, 'x')
+                    tmpVar_y = Variable(key_type, 'y')
+                    if isinstance(key_type.primitive_type, PrimitiveComplexType):
+                        complex_tool_import = Import('pyc_tools_f90', Module('pyc_tools_f90',(),()))
+                        self.add_import(complex_tool_import)
+                        imports_and_macros.append(complex_tool_import)
+                        compare_func = FunctionDef('complex_comparison',
+                                                   [FunctionDefArgument(tmpVar_x), FunctionDefArgument(tmpVar_y)],
+                                                   [FunctionDefResult(Variable(PythonNativeBool(), 'c'))], [])
+                        lt_def = FunctionCall(compare_func, [tmpVar_x, tmpVar_y])
+                    else:
+                        lt_def = PyccelAssociativeParenthesis(PyccelLt(tmpVar_x, tmpVar_y))
+                    imports_and_macros.extend([MacroDefinition('Key', key_type.primitive_type),
+                                   MacroDefinition('Key_KINDLEN(context)', KindSpecification(key_type)),
+                                   MacroDefinition('Key_LT(x,y)', lt_def)])
+                else:
+                    raise NotImplementedError("Support for dicts whose keys define their own < operator is not yet implemented")
+                if isinstance(value_type, FixedSizeNumericType):
+                    imports_and_macros.extend([MacroDefinition('T', value_type.primitive_type),
+                                   MacroDefinition('T_KINDLEN(context)', KindSpecification(value_type))])
+                else:
+                    raise NotImplementedError(f"Support for dictionary values of type {value_type} not yet implemented")
+                imports_and_macros.append(MacroDefinition('Pair', PairType(key_type, value_type)))
+                imports_and_macros.append(MacroDefinition('Map', expr_type))
+                imports_and_macros.append(MacroDefinition('MapIterator', IteratorType(expr_type)))
             else:
                 raise NotImplementedError(f"Unkown gFTL import for type {expr_type}")
 
@@ -1130,6 +1161,24 @@ def _print_PythonSet(self, expr):
             set_type = self._print(expr.class_type)
         return f'{set_type}({list_arg})'
 
+    def _print_PythonDict(self, expr):
+        if len(expr) == 0:
+            list_arg = ''
+            assign = expr.get_direct_user_nodes(lambda a : isinstance(a, Assign))
+            if assign:
+                dict_type = self._print(assign[0].lhs.class_type)
+            else:
+                raise errors.report("Can't use an empty dict without assigning it to a variable as the type cannot be deduced",
+                        severity='fatal', symbol=expr)
+
+        else:
+            class_type = expr.class_type
+            pair_type = self._print(PairType(class_type.key_type, class_type.value_type))
+            args = ', '.join(f'{pair_type}({self._print(k)}, {self._print(v)})' for k,v in expr)
+            list_arg = f'[{args}]'
+            dict_type = self._print(class_type)
+        return f'{dict_type}({list_arg})'
+
     def _print_InhomogeneousTupleVariable(self, expr):
         fs = ', '.join(self._print(f) for f in expr)
         return '[{0}]'.format(fs)
@@ -1988,7 +2037,7 @@ def _print_Allocate(self, expr):
 
             return code
 
-        elif isinstance(class_type, HomogeneousContainerType):
+        elif isinstance(class_type, (HomogeneousContainerType, DictType)):
             return ''
 
         else:
@@ -2006,7 +2055,7 @@ def _print_Deallocate(self, expr):
             Pyccel_del_args = [FunctionCallArgument(var)]
             return self._print(FunctionCall(Pyccel__del, Pyccel_del_args))
 
-        if var.is_alias or isinstance(class_type, (HomogeneousListType, HomogeneousSetType)):
+        if var.is_alias or isinstance(class_type, (HomogeneousListType, HomogeneousSetType, DictType)):
             return ''
         elif isinstance(class_type, (NumpyNDArrayType, HomogeneousTupleType, StringType)):
             var_code = self._print(var)
@@ -2053,9 +2102,16 @@ def _print_HomogeneousListType(self, expr):
     def _print_HomogeneousSetType(self, expr):
         return 'Set_'+self._print(expr.element_type)
 
+    def _print_PairType(self, expr):
+        return 'Pair_'+self._print(expr.key_type)+'__'+self._print(expr.value_type)
+
+    def _print_DictType(self, expr):
+        return 'Map_'+self._print(expr.key_type)+'__'+self._print(expr.value_type)
+
     def _print_IteratorType(self, expr):
         iterable_type = self._print(expr.iterable_type)
         return f"{iterable_type}_Iterator"
+
     def _print_DataType(self, expr):
         return self._print(expr.name)
 

tests/epyccel/test_epyccel_variable_annotations.py

@@ -301,22 +301,32 @@ def homogeneous_list_annotation():
     assert epyc_homogeneous_list_annotation() == homogeneous_list_annotation()
     assert isinstance(epyc_homogeneous_list_annotation(), type(homogeneous_list_annotation()))
 
-def test_dict_int_float(stc_language):
+def test_dict_int_float(language):
     def dict_int_float():
         # Not valid in Python 3.8
         a : dict[int, float] #pylint: disable=unsubscriptable-object,unused-variable
         a = {1:1.0, 2:2.0}
+        return len(a)
 
-    epyc_dict_int_float = epyccel(dict_int_float, language = stc_language)
-    epyc_dict_int_float()
-    dict_int_float()
+    epyc_dict_int_float = epyccel(dict_int_float, language = language)
+    assert epyc_dict_int_float() == dict_int_float()
 
-def test_dict_empty_init(stc_language):
+def test_dict_empty_init(language):
     def dict_empty_init():
         # Not valid in Python 3.8
         a : dict[int, float] #pylint: disable=unsubscriptable-object,unused-variable
         a = {}
+        return len(a)
+
+    epyc_dict_empty_init = epyccel(dict_empty_init, language = language)
+    assert epyc_dict_empty_init() == dict_empty_init()
+
+def test_dict_complex_float(language):
+    def dict_int_float():
+        # Not valid in Python 3.8
+        a : dict[complex, float] #pylint: disable=unsubscriptable-object,unused-variable
+        a = {1j:1.0, -1j:2.0}
+        return len(a)
 
-    epyc_dict_empty_init = epyccel(dict_empty_init, language = stc_language)
-    epyc_dict_empty_init()
-    dict_empty_init()
+    epyc_dict_int_float = epyccel(dict_int_float, language = language)
+    assert epyc_dict_int_float() == dict_int_float()