Add support for set method union (pyccel#2015)

Add support for set method `union`. Fixes pyccel#1753 

**Commit Summary**
- Add a class to represent `SetUnion`
- Remove unused `prefix_code` variable from `CCodePrinter._print_Assign`
- Add C printing for union by adding a function to the extensions
- Add Fortran printing for union by using `copy` and `update`
- Correct error about unknown shape which is irrelevant for objects
whose shape can change
- Add tests for `union`
- Fix AST tree for `DottedName`

Author: EmilyBourne

CHANGELOG.md

@@ -38,6 +38,7 @@ All notable changes to this project will be documented in this file.
 -   #1917 : Add C and Fortran support for set method `add()`.
 -   #1918 : Add support for set method `clear()`.
 -   #1918 : Add support for set method `copy()`.
+-   #1753 : Add support for set method `union()`.
 -   #1936 : Add missing C output for inline decorator example in documentation
 -   #1937 : Optimise `pyccel.ast.basic.PyccelAstNode.substitute` method.
 -   #1544 : Add support for `typing.TypeAlias`.

docs/builtin-functions.md

@@ -111,7 +111,7 @@ Python contains a limited number of builtin functions defined [here](https://doc
 | `remove` | Python-only |
 | `symmetric_difference` | No |
 | `symmetric_difference_update` | No |
-| `union` | No |
+| **`union`** | Yes |
 | `update` | Python-only |
 
 ## Dictionary methods

pyccel/ast/builtin_methods/set_methods.py

@@ -21,6 +21,7 @@
     'SetMethod',
     'SetPop',
     'SetRemove',
+    'SetUnion',
     'SetUpdate'
 )
 
@@ -227,6 +228,36 @@ class SetUpdate(SetMethod):
     """
     __slots__ = ()
     name = 'update'
+    _shape = None
+    _class_type = VoidType()
 
     def __init__(self, set_obj, iterable) -> None:
         super().__init__(set_obj, iterable)
+
+#==============================================================================
+class SetUnion(SetMethod):
+    """
+    Represents a call to the set method .union.
+
+    Represents a call to the set method .union. This method builds a new set
+    by including all elements which appear in at least one of the iterables
+    (the set object and the arguments).
+
+    Parameters
+    ----------
+    set_obj : TypedAstNode
+        The set object which the method is called from.
+    *others : TypedAstNode
+        The iterables which will be combined with this set.
+    """
+    __slots__ = ('_other','_class_type', '_shape')
+    name = 'union'
+
+    def __init__(self, set_obj, *others):
+        self._class_type = set_obj.class_type
+        element_type = self._class_type.element_type
+        for o in others:
+            if element_type != o.class_type.element_type:
+                raise TypeError(f"Argument of type {o.class_type} cannot be used to build set of type {self._class_type}")
+        self._shape = (None,)*self._class_type.rank
+        super().__init__(set_obj, *others)

pyccel/ast/class_defs.py

@@ -6,7 +6,8 @@
 This module contains all types which define a python class which is automatically recognised by pyccel
 """
 
-from pyccel.ast.builtin_methods.set_methods  import SetAdd, SetClear, SetCopy, SetPop, SetRemove, SetDiscard, SetUpdate
+from pyccel.ast.builtin_methods.set_methods  import (SetAdd, SetClear, SetCopy, SetPop, SetRemove,
+                                                     SetDiscard, SetUpdate, SetUnion)
 from pyccel.ast.builtin_methods.list_methods import (ListAppend, ListInsert, ListPop,
                                                      ListClear, ListExtend, ListRemove,
                                                      ListCopy, ListSort)
@@ -163,6 +164,7 @@
             PyccelFunctionDef('discard', func_class = SetDiscard),
             PyccelFunctionDef('pop', func_class = SetPop),
             PyccelFunctionDef('remove', func_class = SetRemove),
+            PyccelFunctionDef('union', func_class = SetUnion),
             PyccelFunctionDef('update', func_class = SetUpdate),
         ])
 

pyccel/ast/core.py

@@ -1970,7 +1970,7 @@ def __init__(self, func, args, current_function=None):
             self._interface = None
             self._funcdef   = func
             self._arguments = tuple(args)
-            self._func_name = func
+            self._func_name = getattr(func, 'name', func)
             super().__init__()
             return
 

pyccel/ast/variable.py

@@ -561,7 +561,7 @@ class DottedName(PyccelAstNode):
     pyccel.stdlib.parallel
     """
     __slots__ = ('_name',)
-    _attribute_nodes = ()
+    _attribute_nodes = ('_name',)
 
     def __new__(cls, *args):
         if len(args) == 1:

pyccel/codegen/printing/ccode.py

@@ -992,7 +992,7 @@ def _print_Import(self, expr):
         else:
             source = expr.source
         if isinstance(source, DottedName):
-            source = source.name[-1]
+            source = source.name[-1].python_value
         else:
             source = self._print(source)
         if source.startswith('stc/') or source in import_header_guard_prefix:
@@ -2228,22 +2228,21 @@ def _print_AugAssign(self, expr):
         return f'{lhs_code} {op}= {rhs_code};\n'
 
     def _print_Assign(self, expr):
-        prefix_code = ''
         lhs = expr.lhs
         rhs = expr.rhs
         if isinstance(rhs, FunctionCall) and isinstance(rhs.class_type, TupleType):
             self._temporary_args = [ObjectAddress(a) for a in lhs]
-            return prefix_code+'{};\n'.format(self._print(rhs))
+            return f'{self._print(rhs)};\n'
         # Inhomogenous tuples are unravelled and therefore do not exist in the c printer
         if isinstance(rhs, (NumpyArray, PythonTuple)):
-            return prefix_code+self.copy_NumpyArray_Data(expr)
+            return self.copy_NumpyArray_Data(expr)
         if isinstance(rhs, (NumpyFull)):
-            return prefix_code+self.arrayFill(expr)
+            return self.arrayFill(expr)
         lhs = self._print(expr.lhs)
         if isinstance(rhs, (PythonList, PythonSet, PythonDict)):
-            return prefix_code+self.init_stc_container(rhs, expr)
+            return self.init_stc_container(rhs, expr)
         rhs = self._print(expr.rhs)
-        return prefix_code+f'{lhs} = {rhs};\n'
+        return f'{lhs} = {rhs};\n'
 
     def _print_AliasAssign(self, expr):
         lhs_var = expr.lhs
@@ -2632,6 +2631,18 @@ def _print_SetCopy(self, expr):
         set_var = self._print(expr.set_variable)
         return f'{var_type}_clone({set_var})'
 
+    def _print_SetUnion(self, expr):
+        assign_base = expr.get_direct_user_nodes(lambda n: isinstance(n, Assign))
+        if not assign_base:
+            errors.report("The result of the union call must be saved into a variable",
+                    severity='error', symbol=expr)
+        class_type = expr.set_variable.class_type
+        var_type = self.get_c_type(class_type)
+        self.add_import(Import('Set_extensions', AsName(VariableTypeAnnotation(class_type), var_type)))
+        set_var = self._print(ObjectAddress(expr.set_variable))
+        args = ', '.join([str(len(expr.args)), *(self._print(ObjectAddress(a)) for a in expr.args)])
+        return f'{var_type}_union({set_var}, {args})'
+
     #=================== MACROS ==================
 
     def _print_MacroShape(self, expr):

pyccel/codegen/printing/fcode.py

@@ -24,6 +24,8 @@
 from pyccel.ast.builtins import PythonTuple, DtypePrecisionToCastFunction
 from pyccel.ast.builtins import PythonBool, PythonList, PythonSet
 
+from pyccel.ast.builtin_methods.set_methods import SetUnion
+
 from pyccel.ast.core import FunctionDef, FunctionDefArgument, FunctionDefResult
 from pyccel.ast.core import SeparatorComment, Comment
 from pyccel.ast.core import ConstructorCall
@@ -262,7 +264,7 @@ def __init__(self, filename, prefix_module = None):
         self._constantImports = {}
         self._current_class    = None
 
-        self._additional_code = None
+        self._additional_code = ''
 
         self.prefix_module = prefix_module
 
@@ -425,8 +427,6 @@ def _handle_inline_func_call(self, expr, assign_lhs = None):
             # Everything before the return node needs handling before the line
             # which calls the inline function is executed
             code = self._print(body)
-            if (not self._additional_code):
-                self._additional_code = ''
             self._additional_code += code
 
             # Collect statements from results to return object
@@ -888,7 +888,7 @@ def _print_Import(self, expr):
 
         source = expr.source
         if isinstance(source, DottedName):
-            source = source.name[-1]
+            source = source.name[-1].python_value
         elif isinstance(source, LiteralString):
             source = source.python_value
         else:
@@ -1276,14 +1276,12 @@ def _print_Constant(self, expr):
     def _print_DottedVariable(self, expr):
         if isinstance(expr.lhs, FunctionCall):
             base = expr.lhs.funcdef.results[0].var
-            if (not self._additional_code):
-                self._additional_code = ''
             var_name = self.scope.get_new_name()
             var = base.clone(var_name)
 
             self.scope.insert_variable(var)
 
-            self._additional_code = self._additional_code + self._print(Assign(var,expr.lhs)) + '\n'
+            self._additional_code += self._print(Assign(var,expr.lhs)) + '\n'
             return self._print(var) + '%' +self._print(expr.name)
         else:
             return self._print(expr.lhs) + '%' +self._print(expr.name)
@@ -1337,6 +1335,31 @@ def _print_SetCopy(self, expr):
         type_name = self._print(expr.class_type)
         return f'{type_name}({var_code})'
 
+    def _print_SetUnion(self, expr):
+        assign_base = expr.get_direct_user_nodes(lambda n: isinstance(n, Assign))
+        var = expr.set_variable
+        if not assign_base:
+            result = self._print(self.scope.get_temporary_variable(var))
+        else:
+            result = self._print(assign_base[0].lhs)
+        expr_type = var.class_type
+        var_code = self._print(expr.set_variable)
+        type_name = self._print(expr_type)
+        self.add_import(self._build_gFTL_extension_module(expr_type))
+        args_insert = []
+        for arg in expr.args:
+            a = self._print(arg)
+            if arg.class_type == expr_type:
+                args_insert.append(f'call {result} % merge({a})\n')
+            else:
+                errors.report(PYCCEL_RESTRICTION_TODO, severity = 'error', symbol = expr)
+        code = f'{result} = {type_name}({var_code})\n' + ''.join(args_insert)
+        if assign_base:
+            return code
+        else:
+            self._additional_code += code
+            return result
+
     #========================== Numpy Elements ===============================#
 
     def _print_NumpySum(self, expr):
@@ -1675,12 +1698,10 @@ def _print_NumpyRand(self, expr):
             errors.report(FORTRAN_ALLOCATABLE_IN_EXPRESSION,
                           symbol=expr, severity='fatal')
 
-        if (not self._additional_code):
-            self._additional_code = ''
         var = self.scope.get_temporary_variable(expr.dtype, memory_handling = 'stack',
                 shape = expr.shape)
 
-        self._additional_code = self._additional_code + self._print(Assign(var,expr)) + '\n'
+        self._additional_code += self._print(Assign(var,expr)) + '\n'
         return self._print(var)
 
     def _print_NumpyRandint(self, expr):
@@ -1978,9 +1999,9 @@ def _print_CodeBlock(self, expr):
         body_stmts = []
         for b in body_exprs :
             line = self._print(b)
-            if (self._additional_code):
+            if self._additional_code:
                 body_stmts.append(self._additional_code)
-                self._additional_code = None
+                self._additional_code = ''
             body_stmts.append(line)
         return ''.join(self._print(b) for b in body_stmts)
 
@@ -2000,7 +2021,7 @@ def _print_NumpyReal(self, expr):
     def _print_Assign(self, expr):
         rhs = expr.rhs
 
-        if isinstance(rhs, FunctionCall):
+        if isinstance(rhs, (FunctionCall, SetUnion)):
             return self._print(rhs)
 
         lhs_code = self._print(expr.lhs)
@@ -3397,17 +3418,13 @@ def _print_FunctionCall(self, expr):
             args = args[1:]
             if isinstance(class_variable, FunctionCall):
                 base = class_variable.funcdef.results[0].var
-                if (not self._additional_code):
-                    self._additional_code = ''
                 var = self.scope.get_temporary_variable(base)
 
-                self._additional_code = self._additional_code + self._print(Assign(var, class_variable)) + '\n'
+                self._additional_code += self._print(Assign(var, class_variable)) + '\n'
                 f_name = f'{self._print(var)} % {f_name}'
             else:
                 f_name = f'{self._print(class_variable)} % {f_name}'
 
-        if (not self._additional_code):
-            self._additional_code = ''
         if parent_assign:
             lhs = parent_assign[0].lhs
             if len(func_results) == 1:

pyccel/codegen/printing/pycode.py

@@ -524,7 +524,8 @@ def _print_EmptyNode(self, expr):
         return ''
 
     def _print_DottedName(self, expr):
-        return '.'.join(self._print(n) for n in expr.name)
+        # A DottedName can only contain LiteralStrings or PyccelSymbols at the printing stage
+        return '.'.join(str(n) for n in expr.name)
 
     def _print_FunctionCall(self, expr):
         func = expr.funcdef

pyccel/parser/semantic.py

@@ -29,7 +29,7 @@
                                  PythonTuple, Lambda, PythonMap)
 
 from pyccel.ast.builtin_methods.list_methods import ListMethod, ListAppend
-from pyccel.ast.builtin_methods.set_methods  import SetMethod, SetAdd
+from pyccel.ast.builtin_methods.set_methods  import SetAdd, SetUnion
 
 from pyccel.ast.core import Comment, CommentBlock, Pass
 from pyccel.ast.core import If, IfSection
@@ -1095,6 +1095,7 @@ def _handle_function(self, expr, func, args, *, is_method = False, use_build_fun
                             new_expr = FunctionCall(func, args)
                         new_expr.set_current_ast(expr.python_ast)
                         pyccel_stage.set_stage('semantic')
+                        new_expr.set_current_user_node(expr.current_user_node)
                         expr = new_expr
                     return getattr(self, annotation_method)(expr)
 
@@ -1566,7 +1567,7 @@ def _assign_lhs_variable(self, lhs, d_var, rhs, new_expressions, is_augassign,ar
                 know_lhs_shape = (lhs.rank == 0) or all(sh is not None for sh in lhs.alloc_shape) \
                         or isinstance(lhs.class_type, StringType)
 
-                if not know_lhs_shape:
+                if isinstance(class_type, (NumpyNDArrayType, HomogeneousTupleType)) and not know_lhs_shape:
                     msg = f"Cannot infer shape of right-hand side for expression {lhs} = {rhs}"
                     errors.report(PYCCEL_RESTRICTION_TODO+'\n'+msg,
                         bounding_box=(self.current_ast_node.lineno, self.current_ast_node.col_offset),
@@ -3158,13 +3159,7 @@ def _visit_Assign(self, expr):
                     symbol=expr, severity='error')
 
         # Checking for the result of _visit_ListExtend
-        if isinstance(rhs, For) or (isinstance(rhs, CodeBlock) and
-            isinstance(rhs.body[0], (ListMethod, SetMethod))):
-            return rhs
-        if isinstance(rhs, ConstructorCall):
-            return rhs
-
-        elif isinstance(rhs, CodeBlock) and len(rhs.body)>1 and isinstance(rhs.body[1], FunctionalFor):
+        if isinstance(rhs, (For, CodeBlock, ConstructorCall)):
             return rhs
 
         elif isinstance(rhs, FunctionCall):
@@ -5032,7 +5027,7 @@ def _build_SetUpdate(self, expr):
         elements of the iterable. If not, it attempts to construct a syntactic `For` 
         loop to iterate over the iterable object and added its elements to the set 
         object. Finally, it passes to a `_visit()` call for semantic parsing.
-    
+
         Parameters
         ----------
         expr : DottedName
@@ -5067,3 +5062,47 @@ def _build_SetUpdate(self, expr):
             pyccel_stage.set_stage('semantic')
             return self._visit(for_obj)
 
+    def _build_SetUnion(self, expr):
+        """
+        Method to navigate the syntactic DottedName node of a `set.union()` call.
+
+        The purpose of this `_build` method is to construct new nodes from a syntactic
+        DottedName node. It creates a SetUnion node if the type of the arguments matches
+        the type of the original set. Otherwise it uses `set.copy` and `set.update` to
+        handle iterators.
+
+        Parameters
+        ----------
+        expr : DottedName
+            The syntactic DottedName node that represent the call to `.union()`.
+
+        Returns
+        -------
+        SetUnion | CodeBlock
+            The nodes describing the union operator.
+        """
+        syntactic_set_obj = expr.name[0]
+        syntactic_args = expr.name[1].args
+        set_obj = self._visit(expr.name[0])
+        args = [self._visit(a.value) for a in expr.name[1].args]
+        class_type = set_obj.class_type
+        if all(a.class_type == class_type for a in args):
+            return SetUnion(set_obj, *args)
+        else:
+            element_type = class_type.element_type
+            if any(a.class_type.element_type != element_type for a in args):
+                errors.report(("Containers containing objects of a different type cannot be used as "
+                               f"arguments to {class_type}.union"),
+                        severity='fatal', symbol=expr)
+
+            lhs = expr.get_user_nodes(Assign)[0].lhs
+            pyccel_stage.set_stage('syntactic')
+            body = [Assign(lhs, DottedName(syntactic_set_obj, FunctionCall('copy', ())),
+                           python_ast = expr.python_ast)]
+            update_calls = [DottedName(lhs, FunctionCall('update', (s_a,))) for s_a in syntactic_args]
+            for c in update_calls:
+                c.set_current_ast(expr.python_ast)
+            body += [Assign(PyccelSymbol('_', is_temp=True), c, python_ast = expr.python_ast)
+                     for c in update_calls]
+            pyccel_stage.set_stage('semantic')
+            return CodeBlock([self._visit(b) for b in body])