Add support for inhomogeneous tuple annotations (pyccel#2066)

Add support for inhomogeneous tuple annotations

Author: EmilyBourne

CHANGELOG.md

@@ -54,6 +54,7 @@ All notable changes to this project will be documented in this file.
 -   #2016 : Add support for translating arithmetic magic methods (methods cannot yet be used from Python).
 -   #1980 : Extend The C support for min and max to more than two variables
 -   #2081 : Add support for multi operator expressions
+-   Add support for inhomogeneous tuple annotations.
 -   \[INTERNALS\] Add abstract class `SetMethod` to handle calls to various set methods.
 -   \[INTERNALS\] Added `container_rank` property to `ast.datatypes.PyccelType` objects.
 -   \[INTERNALS\] Add a `__call__` method to `FunctionDef` to create `FunctionCall` instances.

docs/type_annotations.md

@@ -41,13 +41,29 @@ In general string type hints must be used to provide Pyccel with information abo
 
 ## Tuples
 
-Currently Pyccel supports tuples used locally in functions and in certain cases as arguments, but not as returned objects or module variables. The implementation of the type annotations (including adding the missing support) is in progress. Currently homogeneous tuple type annotations are supported for local variables and function arguments and results (if the tuples contain scalar objects). Internally we handle homogeneous tuples as though they were NumPy arrays. When creating multiple dimensional tuples it is therefore important to ensure that all objects have compatible sizes otherwise they will be handled as inhomogeneous tuples.
+Currently Pyccel supports tuples used locally in functions and in certain cases as arguments, but not as returned objects or module variables. The implementation of the type annotations (including adding the missing support) is in progress.
+
+Tuples can be homogeneous or inhomogeneous. A homogeneous tuple is a tuple whose elements all have the same type and shape. Pyccel translates homogeneous tuples in a similar way to NumPy arrays. When creating multiple dimensional tuples it is therefore important to ensure that all objects have compatible sizes otherwise they will be handled as inhomogeneous tuples. An inhomogeneous tuple describes all other types, but comes with extra restrictions. An inhomogeneous tuple is translated to multiple objects in the target language so it can only be used if the element can be identified during the translation. This means that expressions such as `a[i]` are not possible for inhomogeneous tuples while `a[0]` is valid.
+
+Homogeneous tuple type annotations are supported for local variables and function arguments (if the tuples contain scalar objects).
 
 To declare a homogeneous tuple the syntax is as follows:
 ```python
 a : tuple[int,...] = (1,2,3,4)
 ```
 
+Inhomogeneous tuple type annotations are supported for local variables.
+
+To declare an inhomogeneous tuple the syntax is as follows:
+```python
+a : tuple[int,bool] = (1,False)
+```
+
+It is of course possible to create an inhomogeneous tuple in place of a homogeneous tuple to benefit from code optimisations that can arise from using multiple scalars in place of an array object. This will however imply the same restrictions as any other inhomogeneous tuple. E.g:
+```python
+a : tuple[int, int] = (1,2)
+```
+
 ## Lists
 
 Lists are in the process of being added to Pyccel. Homogeneous lists can be declared in Pyccel using the following syntax:

pyccel/codegen/printing/ccode.py

@@ -1753,10 +1753,10 @@ def _print_Allocate(self, expr):
             if isinstance(variable.class_type, NumpyNDArrayType):
                 #set dtype to the C struct types
                 dtype = self.find_in_ndarray_type_registry(variable.dtype)
-            elif isinstance(variable.class_type, HomogeneousContainerType):
+            elif isinstance(variable.class_type, HomogeneousContainerType) and isinstance(variable.dtype, FixedSizeNumericType):
                 dtype = self.find_in_ndarray_type_registry(numpy_precision_map[(variable.dtype.primitive_type, variable.dtype.precision)])
             else:
-                raise NotImplementedError(f"Don't know how to index {variable.class_type} type")
+                raise NotImplementedError(f"The allocation of the type {variable.class_type} is not yet supported.")
             shape_dtype = self.get_c_type(NumpyInt64Type())
             shape_Assign = "("+ shape_dtype +"[]){" + shape + "}"
             is_view = 'false' if variable.on_heap else 'true'
@@ -1772,9 +1772,9 @@ def _print_Allocate(self, expr):
                     malloc_size = ' * '.join([malloc_size, *(self._print(s) for s in expr.shape)])
                 return f'{var_code} = malloc({malloc_size});\n'
             else:
-                raise NotImplementedError(f"Allocate not implemented for {variable}")
+                raise NotImplementedError(f"Allocate not implemented for {variable.class_type}")
         else:
-            raise NotImplementedError(f"Allocate not implemented for {variable}")
+            raise NotImplementedError(f"Allocate not implemented for {variable.class_type}")
 
     def _print_Deallocate(self, expr):
         if isinstance(expr.variable.class_type, (HomogeneousListType, HomogeneousSetType, DictType)):

pyccel/parser/semantic.py

@@ -72,7 +72,7 @@
 from pyccel.ast.headers import FunctionHeader, MethodHeader, Header
 from pyccel.ast.headers import MacroFunction, MacroVariable
 
-from pyccel.ast.internals import PyccelFunction, Slice, PyccelSymbol
+from pyccel.ast.internals import PyccelFunction, Slice, PyccelSymbol, PyccelArrayShapeElement
 from pyccel.ast.internals import Iterable
 from pyccel.ast.itertoolsext import Product
 
@@ -1565,7 +1565,7 @@ def _assign_lhs_variable(self, lhs, d_var, rhs, new_expressions, is_augassign =
                             for a in args:
                                 if isinstance(a.class_type, InhomogeneousTupleType):
                                     new_args.extend(self.scope.collect_tuple_element(v) for v in a if v.rank>0)
-                                else:
+                                elif a.rank > 0:
                                     new_expressions.append(Allocate(a,
                                         shape=a.alloc_shape, status=status))
                             args = new_args
@@ -1709,6 +1709,26 @@ def _ensure_inferred_type_matches_existing(self, class_type, d_var, var, is_auga
                 tmp_result = PyccelAdd(var, rhs)
                 result_type = tmp_result.class_type
                 raise_error = var.class_type != result_type
+            elif isinstance(var.class_type, InhomogeneousTupleType) and \
+                    isinstance(class_type, HomogeneousTupleType):
+                if d_var['shape'][0] == var.shape[0]:
+                    rhs_elem = self.scope.collect_tuple_element(var[0])
+                    self._ensure_inferred_type_matches_existing(class_type.element_type,
+                            self._infer_type(rhs_elem), rhs_elem, is_augassign, new_expressions, rhs)
+                    raise_error = False
+                else:
+                    raise_error = True
+            elif isinstance(var.class_type, InhomogeneousTupleType) and \
+                    isinstance(class_type, InhomogeneousTupleType):
+                for i, element_type in enumerate(class_type):
+                    rhs_elem = self.scope.collect_tuple_element(var[i])
+                    self._ensure_inferred_type_matches_existing(element_type,
+                            self._infer_type(rhs_elem), rhs_elem, is_augassign, new_expressions, rhs)
+                raise_error = False
+            elif isinstance(var.class_type, HomogeneousTupleType) and \
+                    isinstance(class_type, InhomogeneousTupleType):
+                # TODO: Remove isinstance(rhs, Variable) condition when tuples are saved like lists
+                raise_error = any(a != var.class_type.element_type for a in class_type) or not isinstance(rhs, Variable)
             else:
                 raise_error = True
 
@@ -2027,7 +2047,8 @@ def _get_indexed_type(self, base, args, expr):
             else:
                 raise errors.report(f"Unknown annotation base {base}\n"+PYCCEL_RESTRICTION_TODO,
                         severity='fatal', symbol=expr)
-            if (len(args) == 2 and args[1] is LiteralEllipsis()) or len(args) == 1:
+            if (len(args) == 2 and args[1] is LiteralEllipsis()) or \
+                    (len(args) == 1 and dtype_cls is not PythonTupleFunction):
                 syntactic_annotation = self._convert_syntactic_object_to_type_annotation(args[0])
                 internal_datatypes = self._visit(syntactic_annotation)
                 if dtype_cls in type_container:
@@ -2046,6 +2067,13 @@ def _get_indexed_type(self, base, args, expr):
                 type_annotations = [VariableTypeAnnotation(DictType(k.class_type, v.class_type)) \
                                     for k,v in zip(key_types.type_list, val_types.type_list)]
                 return UnionTypeAnnotation(*type_annotations)
+            elif dtype_cls is PythonTupleFunction:
+                syntactic_annotations = [self._convert_syntactic_object_to_type_annotation(a) for a in args]
+                types = [self._visit(a).type_list for a in syntactic_annotations]
+                internal_datatypes = list(product(*types))
+                type_annotations = [VariableTypeAnnotation(InhomogeneousTupleType(*[ui.class_type for ui in u]), True)
+                                    for u in internal_datatypes]
+                return UnionTypeAnnotation(*type_annotations)
             else:
                 raise errors.report("Cannot handle non-homogenous type index\n"+PYCCEL_RESTRICTION_TODO,
                         severity='fatal', symbol=expr)
@@ -2808,12 +2836,21 @@ def _visit_AnnotatedPyccelSymbol(self, expr):
             elif isinstance(t, VariableTypeAnnotation):
                 class_type = t.class_type
                 cls_base = self.scope.find(str(class_type), 'classes') or get_cls_base(class_type)
+                shape = len(class_type) if isinstance(class_type, InhomogeneousTupleType) else None
                 v = var_class(class_type, name, cls_base = cls_base,
-                        shape = None,
+                        shape = shape,
                         is_const = t.is_const, is_optional = False,
                         memory_handling = array_memory_handling if class_type.rank > 0 else 'stack',
                         **kwargs)
                 possible_args.append(v)
+                if isinstance(class_type, InhomogeneousTupleType):
+                    for i, t in enumerate(class_type):
+                        pyccel_stage.set_stage('syntactic')
+                        syntactic_elem = AnnotatedPyccelSymbol(self.scope.get_new_name( f'{name}_{i}'),
+                                                annotation = UnionTypeAnnotation(VariableTypeAnnotation(t)))
+                        pyccel_stage.set_stage('semantic')
+                        elem = self._visit(syntactic_elem)
+                        self.scope.insert_symbolic_alias(IndexedElement(v, i), elem[0])
             else:
                 errors.report(PYCCEL_RESTRICTION_TODO + '\nUnrecoginsed type annotation',
                         severity='fatal', symbol=expr)
@@ -3471,7 +3508,8 @@ def _visit_Assign(self, expr):
                 if isinstance(l.class_type, InhomogeneousTupleType) \
                         and not isinstance(r, (FunctionCall, PyccelFunction)):
                     new_lhs.extend(self.scope.collect_tuple_element(v) for v in l)
-                    new_rhs.extend(self.scope.collect_tuple_element(v) for v in r)
+                    new_rhs.extend(self.scope.collect_tuple_element(r[i]) \
+                            for i in range(len(l.class_type)))
                     # Repeat step to handle tuples of tuples of etc.
                     unravelling = True
                 elif isinstance(l, Variable) and isinstance(l.class_type, InhomogeneousTupleType):
@@ -3509,7 +3547,8 @@ def _visit_Assign(self, expr):
                             bounding_box=(self.current_ast_node.lineno, self.current_ast_node.col_offset),
                             symbol=li, severity='error')
             else:
-                if getattr(l, 'is_const', False) and (not isinstance(expr.lhs, AnnotatedPyccelSymbol) or len(l.get_all_user_nodes()) > 0):
+                if getattr(l, 'is_const', False) and (not isinstance(expr.lhs, AnnotatedPyccelSymbol) or \
+                        any(not isinstance(u, (Allocate, PyccelArrayShapeElement)) for u in l.get_all_user_nodes())):
                     # If constant and not the initialising declaration of a constant variable
                     errors.report("Cannot modify 'const' variable",
                         bounding_box=(self.current_ast_node.lineno, self.current_ast_node.col_offset),

tests/epyccel/test_epyccel_variable_annotations.py

@@ -324,13 +324,77 @@ def 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 : dict[complex, float] #pylint: disable=unsubscriptable-object
         a = {1j:1.0, -1j:2.0}
         return len(a)
 
     epyc_dict_int_float = epyccel(dict_int_float, language = language)
     assert epyc_dict_int_float() == dict_int_float()
 
+def test_inhomogeneous_tuple_annotation_1(language):
+    def inhomogeneous_tuple_annotation():
+        a : tuple[int, bool] = (1, True) #pylint: disable=unsubscriptable-object
+        return a[0], a[1]
+
+    epyc_inhomogeneous_tuple_annotation = epyccel(inhomogeneous_tuple_annotation, language = language)
+    assert epyc_inhomogeneous_tuple_annotation() == inhomogeneous_tuple_annotation()
+
+def test_inhomogeneous_tuple_annotation_2(language):
+    def inhomogeneous_tuple_annotation():
+        a : tuple[int] = (1,) #pylint: disable=unsubscriptable-object
+        return a[0]
+
+    epyc_inhomogeneous_tuple_annotation = epyccel(inhomogeneous_tuple_annotation, language = language)
+    assert epyc_inhomogeneous_tuple_annotation() == inhomogeneous_tuple_annotation()
+
+def test_inhomogeneous_tuple_annotation_3(language):
+    def inhomogeneous_tuple_annotation():
+        a : tuple[int,int,int] = (1,2,3) #pylint: disable=unsubscriptable-object
+        return a[0], a[1], a[2]
+
+    epyc_inhomogeneous_tuple_annotation = epyccel(inhomogeneous_tuple_annotation, language = language)
+    assert epyc_inhomogeneous_tuple_annotation() == inhomogeneous_tuple_annotation()
+
+def test_inhomogeneous_tuple_annotation_4(language):
+    def inhomogeneous_tuple_annotation():
+        a : tuple[tuple[float,bool],tuple[int,complex]] = ((1.0, False), (1,2+3j)) #pylint: disable=unsubscriptable-object
+        return a[0][0], a[0][1], a[1][0], a[1][1]
+
+    epyc_inhomogeneous_tuple_annotation = epyccel(inhomogeneous_tuple_annotation, language = language)
+    assert epyc_inhomogeneous_tuple_annotation() == inhomogeneous_tuple_annotation()
+
+def test_inhomogeneous_tuple_annotation_5(language):
+    def inhomogeneous_tuple_annotation():
+        a : tuple[tuple[int, float]] = ((1,0.2),) #pylint: disable=unsubscriptable-object
+        return a[0][0], a[0][1]
+
+    epyc_inhomogeneous_tuple_annotation = epyccel(inhomogeneous_tuple_annotation, language = language)
+    assert epyc_inhomogeneous_tuple_annotation() == inhomogeneous_tuple_annotation()
+
+def test_inhomogeneous_tuple_annotation_6(language):
+    def inhomogeneous_tuple_annotation():
+        a : tuple[tuple[tuple[int, float]]] = (((1,0.2),),) #pylint: disable=unsubscriptable-object
+        return a[0][0][0], a[0][0][1]
+
+    epyc_inhomogeneous_tuple_annotation = epyccel(inhomogeneous_tuple_annotation, language = language)
+    assert epyc_inhomogeneous_tuple_annotation() == inhomogeneous_tuple_annotation()
+
+def test_inhomogeneous_tuple_annotation_7(language):
+    def inhomogeneous_tuple_annotation():
+        a : tuple[tuple[tuple[int, float]], int] = (((1,0.2),),1) #pylint: disable=unsubscriptable-object
+        return a[0][0][0], a[0][0][1], a[1]
+
+    epyc_inhomogeneous_tuple_annotation = epyccel(inhomogeneous_tuple_annotation, language = language)
+    assert epyc_inhomogeneous_tuple_annotation() == inhomogeneous_tuple_annotation()
+
+def test_inhomogeneous_tuple_annotation_8(language):
+    def inhomogeneous_tuple_annotation():
+        a : tuple[tuple[tuple[tuple[int, float]], int]] = ((((1,0.2),),1),) #pylint: disable=unsubscriptable-object
+        return a[0][0][0][0], a[0][0][0][1], a[0][1]
+
+    epyc_inhomogeneous_tuple_annotation = epyccel(inhomogeneous_tuple_annotation, language = language)
+    assert epyc_inhomogeneous_tuple_annotation() == inhomogeneous_tuple_annotation()
+
 @pytest.mark.parametrize( 'language', (
         pytest.param("fortran", marks = pytest.mark.fortran),
         pytest.param("c", marks = [