Add support for returning lists, sets and homogeneous tuples from C (pyccel#2070)

Generalise `_extract_HomogeneousTupleType_FunctionDefResult` to
`_extract_HomogeneousContainerType_FunctionDefResult` to add support for
returning lists, sets and homogeneous tuples from C. Fixes pyccel#1664 and
fixes pyccel#1665
Improve docs

Author: EmilyBourne

CHANGELOG.md

@@ -33,9 +33,11 @@ All notable changes to this project will be documented in this file.
 -   #1874 : Add C and Fortran support for the `len()` function for the `list` container.
 -   #1875 : Add C and Fortran support for the `len()` function for the `set` container.
 -   #1908 : Add C and Fortran support for the `len()` function for the `dict` container.
+-   #1665 : Add C support for returning lists from functions.
 -   #1689 : Add C and Fortran support for list method `append()`.
 -   #1876 : Add C support for indexing lists.
 -   #1690 : Add C support for list method `pop()`.
+-   #1664 : Add C support for returning sets from functions.
 -   #2023 : Add support for iterating over a `set`.
 -   #1877 : Add C and Fortran Support for set method `pop()`.
 -   #1917 : Add C and Fortran support for set method `add()`.
@@ -51,6 +53,7 @@ All notable changes to this project will be documented in this file.
 -   Add a warning about containers in lists.
 -   #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
 -   \[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,33 @@ 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 (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. 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.
 
 To declare a homogeneous tuple the syntax is as follows:
 ```python
 a : tuple[int,...] = (1,2,3,4)
 ```
 
+## Lists
+
+Lists are in the process of being added to Pyccel. Homogeneous lists can be declared in Pyccel using the following syntax:
+```python
+a : list[int] = [1, 2]
+b : list[bool] = [False, True]
+c : list[float] = []
+```
+So far lists can be declared as local variables or as results of functions.
+
+## Sets
+
+Sets are in the process of being added to Pyccel. Homogeneous sets can be declared in Pyccel using the following syntax:
+```python
+a : set[int] = {1, 2}
+b : set[bool] = {False, True}
+c : set[float] = {}
+```
+So far sets can be declared as local variables or as results of functions.
+
 ## Dictionaries
 
 Dictionaries are in the process of being added to Pyccel. They cannot yet be used effectively however the type annotations are already supported.

pyccel/ast/cwrapper.py

@@ -1105,7 +1105,7 @@ def C_to_Python(c_object):
                         arguments = [FunctionDefArgument(Variable(PyccelPyObject(), name='l', memory_handling='alias')),
                                      FunctionDefArgument(Variable(PythonNativeInt(), name='i')),
                                      FunctionDefArgument(Variable(PyccelPyObject(), name='new_item', memory_handling='alias'))],
-                        results = [])
+                        results = [FunctionDefResult(Variable(CNativeInt(), 'i'))])
 
 #-------------------------------------------------------------------
 #                         Tuple functions
@@ -1136,6 +1136,31 @@ def C_to_Python(c_object):
                                      FunctionDefArgument(Variable(PythonNativeInt(), name='i'))],
                         results = [FunctionDefResult(Variable(PyccelPyObject(), name='o', memory_handling='alias'))])
 
+# https://docs.python.org/3/c-api/tuple.html#c.PyTuple_SetItem
+PyTuple_SetItem = FunctionDef(name = 'PyTuple_SetItem',
+                        body = [],
+                        arguments = [FunctionDefArgument(Variable(PyccelPyObject(), name='l', memory_handling='alias')),
+                                     FunctionDefArgument(Variable(PythonNativeInt(), name='i')),
+                                     FunctionDefArgument(Variable(PyccelPyObject(), name='new_item', memory_handling='alias'))],
+                        results = [FunctionDefResult(Variable(CNativeInt(), 'i'))])
+
+#-------------------------------------------------------------------
+#                          Set functions
+#-------------------------------------------------------------------
+
+# https://docs.python.org/3/c-api/set.html#c.PySet_New
+PySet_New = FunctionDef(name = 'PySet_New',
+                    arguments = [FunctionDefArgument(Variable(PyccelPyObject(), 'iterable', memory_handling='alias'), value = Nil())],
+                    results = [FunctionDefResult(Variable(PyccelPyObject(), 'set', memory_handling='alias'))],
+                    body = [])
+
+# https://docs.python.org/3/c-api/set.html#c.PySet_Add
+PySet_Add = FunctionDef(name = 'PySet_Add',
+                    arguments = [FunctionDefArgument(Variable(PyccelPyObject(), 'set', memory_handling='alias')),
+                                 FunctionDefArgument(Variable(PyccelPyObject(), 'key', memory_handling='alias'))],
+                    results = [FunctionDefResult(Variable(PythonNativeInt(), 'i'))],
+                    body = [])
+
 
 # Functions definitions are defined in pyccel/stdlib/cwrapper/cwrapper.c
 check_type_registry = {

pyccel/codegen/printing/cwrappercode.py

@@ -267,6 +267,7 @@ def _print_ModuleHeader(self, expr):
         imports  = [*module_imports, *mod.imports]
         for i in imports:
             self.add_import(i)
+        self.invalidate_stc_headers(imports)
         imports  = ''.join(self._print(i) for i in imports)
 
         function_signatures = ''.join(self.function_signature(f, print_arg_names = False) + ';\n' for f in mod.external_funcs)
@@ -373,6 +374,7 @@ def _print_PyModule(self, expr):
 
         pymod_name = f'{expr.name}_wrapper'
         imports = [Import(pymod_name, Module(pymod_name,(),())), *self._additional_imports.values()]
+        self.invalidate_stc_headers(imports)
         imports  = ''.join(self._print(i) for i in imports)
 
         self.exit_scope()

pyccel/codegen/wrapper/c_to_python_wrapper.py

@@ -11,7 +11,8 @@
 from pyccel.ast.bind_c        import BindCFunctionDef, BindCPointer, BindCFunctionDefArgument
 from pyccel.ast.bind_c        import BindCModule, BindCVariable, BindCFunctionDefResult
 from pyccel.ast.bind_c        import BindCClassDef, BindCClassProperty
-from pyccel.ast.builtins      import PythonTuple, PythonRange
+from pyccel.ast.builtins      import PythonTuple, PythonRange, PythonLen
+from pyccel.ast.builtin_methods.set_methods import SetPop
 from pyccel.ast.class_defs    import StackArrayClass
 from pyccel.ast.core          import Interface, If, IfSection, Return, FunctionCall
 from pyccel.ast.core          import FunctionDef, FunctionDefArgument, FunctionDefResult
@@ -29,10 +30,13 @@
 from pyccel.ast.cwrapper      import PyList_New, PyList_Append, PyList_GetItem, PyList_SetItem
 from pyccel.ast.cwrapper      import PyccelPyTypeObject, PyCapsule_New, PyCapsule_Import
 from pyccel.ast.cwrapper      import PySys_GetObject, PyUnicode_FromString, PyGetSetDefElement
-from pyccel.ast.cwrapper      import PyTuple_Size, PyTuple_Check, PyTuple_GetItem
+from pyccel.ast.cwrapper      import PyTuple_Size, PyTuple_Check, PyTuple_New
+from pyccel.ast.cwrapper      import PyTuple_GetItem, PyTuple_SetItem
+from pyccel.ast.cwrapper      import PySet_New, PySet_Add
 from pyccel.ast.c_concepts    import ObjectAddress, PointerCast, CStackArray, CNativeInt
 from pyccel.ast.datatypes     import VoidType, PythonNativeInt, CustomDataType, DataTypeFactory
 from pyccel.ast.datatypes     import FixedSizeNumericType, HomogeneousTupleType, PythonNativeBool
+from pyccel.ast.datatypes     import HomogeneousSetType, HomogeneousListType
 from pyccel.ast.datatypes     import TupleType
 from pyccel.ast.literals      import Nil, LiteralTrue, LiteralString, LiteralInteger
 from pyccel.ast.literals      import LiteralFalse, convert_to_literal
@@ -761,6 +765,7 @@ def _build_module_import_function(self, expr):
 
         ok_code = LiteralInteger(0, dtype=CNativeInt())
         error_code = PyccelUnarySub(LiteralInteger(1, dtype=CNativeInt()))
+        self._error_exit_code = error_code
 
         # Create variables to temporarily modify the Python path so the file will be discovered
         current_path = func_scope.get_temporary_variable(PyccelPyObject(), 'current_path', memory_handling='alias')
@@ -769,14 +774,19 @@ def _build_module_import_function(self, expr):
         body = [AliasAssign(current_path, PySys_GetObject(LiteralString("path"))),
                 AliasAssign(stash_path, PyList_GetItem(current_path, LiteralInteger(0, dtype=CNativeInt()))),
                 Py_INCREF(stash_path),
-                PyList_SetItem(current_path, LiteralInteger(0, dtype=CNativeInt()),
-                               PyUnicode_FromString(LiteralString(self._file_location))),
+                If(IfSection(PyccelEq(PyList_SetItem(current_path, LiteralInteger(0, dtype=CNativeInt()),
+                                                PyUnicode_FromString(LiteralString(self._file_location))),
+                                      PyccelUnarySub(LiteralInteger(1))),
+                             [Return([self._error_exit_code])])),
                 AliasAssign(API_var, PyCapsule_Import(self.scope.get_python_name(mod_name))),
-                PyList_SetItem(current_path, LiteralInteger(0, dtype=CNativeInt()), stash_path),
+                If(IfSection(PyccelEq(PyList_SetItem(current_path, LiteralInteger(0, dtype=CNativeInt()), stash_path),
+                                      PyccelUnarySub(LiteralInteger(1))),
+                             [Return([self._error_exit_code])])),
                 Return([IfTernaryOperator(PyccelIsNot(API_var, Nil()), ok_code, error_code)])]
 
         result = func_scope.get_temporary_variable(CNativeInt())
         self.exit_scope()
+        self._error_exit_code = Nil()
         import_func = FunctionDef(func_name, (), (FunctionDefResult(result),), body, is_static=True, scope = func_scope)
 
         return API_var, import_func
@@ -2470,7 +2480,7 @@ def _extract_FixedSizeType_FunctionDefResult(self, orig_var, is_bind_c, funcdef)
         dict
             A dictionary describing the objects necessary to collect the result.
         """
-        name = orig_var.name
+        name = getattr(orig_var, 'name', 'tmp')
         py_res = self.get_new_PyObject(f'{name}_obj', orig_var.dtype)
         c_res = Variable(orig_var.class_type, self.scope.get_new_name(name))
         self.scope.insert_variable(c_res)
@@ -2571,13 +2581,14 @@ def _extract_InhomogeneousTupleType_FunctionDefResult(self, orig_var, is_bind_c,
         body.extend(Py_DECREF(r) for r in py_result_vars)
         return {'c_results': c_result_vars, 'py_result': py_res, 'body': body, 'setup': setup}
 
-    def _extract_HomogeneousTupleType_FunctionDefResult(self, orig_var, is_bind_c, funcdef):
+    def _extract_HomogeneousContainerType_FunctionDefResult(self, orig_var, is_bind_c, funcdef):
         """
-        Get the code which translates a `Variable` containing a homogeneous tuple to a PyObject.
+        Get the code which translates a `Variable` containing a homogeneous container to a PyObject.
 
-        Get the code which translates a `Variable` containing a homogeneous tuple to a PyObject.
-        For now this function does not handle homogeneous tuples but merely allows multiple
-        returns of the same type to be outputted neatly.
+        Get the code which translates a `Variable` containing a homogeneous container to a PyObject.
+        This function handles lists, sets, and tuples.
+        The current implementation for tuples is not working correctly as Pyccel does not handle
+        function calls to functions returning tuples correctly.
 
         Parameters
         ----------
@@ -2596,6 +2607,46 @@ def _extract_HomogeneousTupleType_FunctionDefResult(self, orig_var, is_bind_c, f
         """
         if isinstance(orig_var, PythonTuple):
             return self._extract_InhomogeneousTupleType_FunctionDefResult(orig_var, is_bind_c, funcdef)
+
+        if is_bind_c:
+            raise NotImplementedError("Support for returning sets from Fortran code is not yet available")
+
+        name = getattr(orig_var, 'name', 'tmp')
+        py_res = self.get_new_PyObject(f'{name}_obj', orig_var.dtype)
+        c_res = orig_var.clone(self.scope.get_new_name(name), is_argument = False)
+        loop_size = Variable(PythonNativeInt(), self.scope.get_new_name(f'{name}_size'))
+        idx = Variable(PythonNativeInt(), self.scope.get_new_name())
+        self.scope.insert_variable(c_res)
+        self.scope.insert_variable(loop_size)
+        self.scope.insert_variable(idx)
+
+        for_scope = self.scope.create_new_loop_scope()
+        self.scope = for_scope
+        element_extraction = self._extract_FunctionDefResult(IndexedElement(orig_var, idx), is_bind_c, funcdef)
+        self.exit_scope()
+
+        class_type = orig_var.class_type
+        if isinstance(class_type, HomogeneousSetType):
+            element = SetPop(c_res)
+            elem_set = PySet_Add(py_res, element_extraction['py_result'])
+            init = PySet_New()
+        elif isinstance(class_type, HomogeneousListType):
+            element = IndexedElement(c_res, idx)
+            elem_set = PyList_SetItem(py_res, idx, element_extraction['py_result'])
+            init = PyList_New(loop_size)
+        elif isinstance(class_type, HomogeneousTupleType):
+            element = IndexedElement(c_res, idx)
+            elem_set = PyTuple_SetItem(py_res, idx, element_extraction['py_result'])
+            init = PyTuple_New(loop_size)
         else:
-            return errors.report(f"Wrapping function results is not implemented for type {orig_var.class_type}. " + PYCCEL_RESTRICTION_TODO, symbol=orig_var,
-                severity='fatal')
+            raise NotImplementedError(f"Don't know how to return an object of type {class_type}")
+
+        for_body = [Assign(element_extraction['c_results'][0], element),
+                *element_extraction['body'],
+                If(IfSection(PyccelEq(elem_set, PyccelUnarySub(LiteralInteger(1))),
+                                         [Return([self._error_exit_code])]))]
+        body = [Assign(loop_size, PythonLen(c_res)),
+                AliasAssign(py_res, init),
+                For((idx,), PythonRange(loop_size), for_body, for_scope)]
+
+        return {'c_results': [c_res], 'py_result': py_res, 'body': body}

tests/epyccel/test_epyccel_lists.py

@@ -769,3 +769,15 @@ def list_ptr():
     python_result = list_ptr()
     assert isinstance(python_result, type(pyccel_result))
     assert python_result == pyccel_result
+
+def test_list_return(stc_language):
+    def list_return():
+        a = [1,2,3,4,5]
+        return a
+
+    epyccel_func = epyccel(list_return, language = stc_language)
+    pyccel_result = epyccel_func()
+    python_result = list_return()
+    assert python_result == pyccel_result
+    assert isinstance(python_result, type(pyccel_result))
+    assert isinstance(python_result.pop(), type(pyccel_result.pop()))

tests/epyccel/test_epyccel_sets.py

@@ -562,3 +562,17 @@ def set_iter_prod():
     python_result = set_iter_prod()
     assert python_result == pyccel_result
     assert isinstance(python_result, type(pyccel_result))
+
+def test_set_return(stc_language):
+    def set_return():
+        a = {1,2,3,4,5}
+        b = {4,5,6}
+        c = a.union(b) # Use union to avoid #2084
+        return c
+
+    epyccel_func = epyccel(set_return, language = stc_language)
+    pyccel_result = epyccel_func()
+    python_result = set_return()
+    assert python_result == pyccel_result
+    assert isinstance(python_result, type(pyccel_result))
+    assert isinstance(python_result.pop(), type(pyccel_result.pop()))

tests/epyccel/test_tuples.py

@@ -185,3 +185,13 @@ def my_tuple(a : 'tuple[tuple[int,...],...]'):
         epyc_func = epyccel(my_tuple, language=language)
 
         assert my_tuple(tuple_arg) == epyc_func(tuple_arg)
+
+@pytest.mark.xfail(reason="Returning tuples from functions requires a reorganisation of the return system. See #337")
+def test_homogeneous_tuples_result(stc_language):
+    def my_tuple() -> 'tuple[int, ...]':
+        a = (1,2,3,4,5)
+        return a
+
+    epyc_func = epyccel(my_tuple, language=stc_language)
+
+    assert my_tuple() == epyc_func()