Add support for in operator (pyccel#2017)

Add support for `in` operator for `list`, `set` and `dict` containers
for C Fortran and Python, as well as for classes.

**Commit Summary**
- Add `PyccelIn` class to describe the `in` operator
- Add C support for `in` operator for list/set/dict
- Add `#define i_use_cmp` to list declaration in C to activate sorting
and searching functions
- Add `_define_gFTL_element` function in Fortran code printer to define
types and operators for types inside container types
- Reduce duplication in `_build_gFTL_module`
- Add Fortran support for `in` operator for list/set/dict
- Add Python support for `in` operator for list/set/dict
- In semantic stage handle `PyccelIn` for built-in types and class
`__contains__` magic method
- Add tests for `in` operator for lists/sets/dicts
- Add tests for `in` operator for classes
- Use `python_only_language` to allow some of the dictionary tests to be
activated for other languages

Author: EmilyBourne

CHANGELOG.md

@@ -27,6 +27,7 @@ All notable changes to this project will be documented in this file.
 -   #1657 : Add the appropriate Fortran language equivalent for declaring a Python `list` container using the gFTL library.
 -   #1658 : Add the appropriate Fortran language equivalent for declaring a Python `set` container using the gFTL library.
 -   #1944 : Add the appropriate Fortran language equivalent for declaring a Python `dict` container using the gFTL library.
+-   #2009 : Add support for `in` operator for `list`, `set`, `dict` and class containers.
 -   #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.

pyccel/ast/operators.py

@@ -56,6 +56,7 @@
     'Relational',
     'PyccelIs',
     'PyccelIsNot',
+    'PyccelIn',
     'IfTernaryOperator'
 )
 
@@ -1285,6 +1286,46 @@ def eval(self):
 
 #==============================================================================
 
+class PyccelIn(PyccelBooleanOperator):
+    """
+    Represents an `in` expression in the code.
+
+    Represents an `in` expression in the code.
+
+    Parameters
+    ----------
+    element : TypedAstNode
+        The first argument passed to the operator.
+
+    container : TypedAstNode
+        The first argument passed to the operator.
+    """
+    __slots__ = ()
+    _precedence = 7
+
+    def __init__(self, element, container):
+        super().__init__(element, container)
+
+    @property
+    def element(self):
+        """
+        First operator argument.
+
+        First operator argument.
+        """
+        return self._args[0]
+
+    @property
+    def container(self):
+        """
+        Second operator argument.
+
+        Second operator argument.
+        """
+        return self._args[1]
+
+#==============================================================================
+
 class IfTernaryOperator(PyccelOperator):
     """
     Represent a ternary conditional operator in the code.

pyccel/codegen/printing/ccode.py

@@ -972,6 +972,20 @@ def _print_PyccelNot(self, expr):
         a = self._print(expr.args[0])
         return '!{}'.format(a)
 
+    def _print_PyccelIn(self, expr):
+        container_type = expr.container.class_type
+        element = self._print(expr.element)
+        container = self._print(ObjectAddress(expr.container))
+        c_type = self.get_c_type(expr.container.class_type)
+        if isinstance(container_type, (HomogeneousSetType, DictType)):
+            return f'{c_type}_contains({container}, {element})'
+        elif isinstance(container_type, HomogeneousListType):
+            return f'{c_type}_find({container}, {element}).ref != {c_type}_end({container}).ref'
+        else:
+            raise errors.report(PYCCEL_RESTRICTION_TODO,
+                    symbol = expr,
+                    severity='fatal')
+
     def _print_PyccelMod(self, expr):
         self.add_import(c_imports['math'])
         self.add_import(c_imports['pyc_math_c'])
@@ -1018,16 +1032,19 @@ def _print_Import(self, expr):
         if source.startswith('stc/') or source in import_header_guard_prefix:
             code = ''
             for t in expr.target:
-                dtype = t.object.class_type
+                class_type = t.object.class_type
                 container_type = t.local_alias
-                if isinstance(dtype, DictType):
-                    container_key_key = self.get_c_type(dtype.key_type)
-                    container_val_key = self.get_c_type(dtype.value_type)
+                if isinstance(class_type, DictType):
+                    container_key_key = self.get_c_type(class_type.key_type)
+                    container_val_key = self.get_c_type(class_type.value_type)
                     container_key = f'{container_key_key}_{container_val_key}'
                     element_decl = f'#define i_key {container_key_key}\n#define i_val {container_val_key}\n'
                 else:
-                    container_key = self.get_c_type(dtype.element_type)
+                    container_key = self.get_c_type(class_type.element_type)
                     element_decl = f'#define i_key {container_key}\n'
+                if isinstance(class_type, HomogeneousListType) and isinstance(class_type.element_type, FixedSizeNumericType) \
+                        and not isinstance(class_type.element_type.primitive_type, PrimitiveComplexType):
+                    element_decl += '#define i_use_cmp\n'
                 header_guard_prefix = import_header_guard_prefix.get(source, '')
                 header_guard = f'{header_guard_prefix}_{container_type.upper()}'
                 code += ''.join((f'#ifndef {header_guard}\n',

pyccel/codegen/printing/fcode.py

@@ -69,7 +69,7 @@
 
 from pyccel.ast.numpytypes import NumpyNDArrayType
 
-from pyccel.ast.operators import PyccelAdd, PyccelMul, PyccelMinus, PyccelAnd
+from pyccel.ast.operators import PyccelAdd, PyccelMul, PyccelMinus, PyccelAnd, PyccelEq
 from pyccel.ast.operators import PyccelMod, PyccelNot, PyccelAssociativeParenthesis
 from pyccel.ast.operators import PyccelUnarySub, PyccelLt, PyccelGt, IfTernaryOperator
 
@@ -524,12 +524,12 @@ def _calculate_class_names(self, expr):
                         scope.rename_function(f, suggested_name)
                     f.cls_name = scope.get_new_name(f'{name}_{f.name}')
 
-    def _get_comparison_operator(self, element_type, imports_and_macros):
+    def _define_gFTL_element(self, element_type, imports_and_macros, element_name):
         """
-        Get an AST node describing a comparison operator between two objects of the same type.
+        Get lists of nodes describing comparison operators between two objects of the same type.
 
-        Get an AST node describing a comparison operator between two objects of the same type.
-        This is useful when defining gFTL modules.
+        Get lists of nodes describing a comparison operators between two objects of the same type.
+        This is necessary when defining gFTL modules.
 
         Parameters
         ----------
@@ -539,25 +539,48 @@ def _get_comparison_operator(self, element_type, imports_and_macros):
         imports_and_macros : list
             A list of imports or macros for the gFTL module.
 
+        element_name : str
+            The name of the element whose properties are being specified.
+
         Returns
         -------
-        PyccelAstNode
-            An AST node describing a comparison operator.
+        defs : list[MacroDefinition]
+            A list of nodes describing the macros defining comparison operator.
+        undefs : list[MacroUndef]
+            A list of nodes which undefine the macros.
         """
-        tmpVar_x = Variable(element_type, 'x')
-        tmpVar_y = Variable(element_type, 'y')
-        if isinstance(element_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 = compare_func(tmpVar_x, tmpVar_y)
+        if isinstance(element_type, FixedSizeNumericType):
+            tmpVar_x = Variable(element_type, 'x')
+            tmpVar_y = Variable(element_type, 'y')
+            if isinstance(element_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 = compare_func(tmpVar_x, tmpVar_y)
+            else:
+                lt_def = PyccelAssociativeParenthesis(PyccelLt(tmpVar_x, tmpVar_y))
+
+            defs = [MacroDefinition(element_name, element_type.primitive_type),
+                    MacroDefinition(f'{element_name}_KINDLEN(context)', KindSpecification(element_type)),
+                    MacroDefinition(f'{element_name}_LT(x,y)', lt_def),
+                    MacroDefinition(f'{element_name}_EQ(x,y)', PyccelAssociativeParenthesis(PyccelEq(tmpVar_x, tmpVar_y)))]
+            undefs = [MacroUndef(element_name),
+                      MacroUndef(f'{element_name}_KINDLEN'),
+                      MacroUndef(f'{element_name}_LT'),
+                      MacroUndef(f'{element_name}_EQ')]
         else:
-            lt_def = PyccelAssociativeParenthesis(PyccelLt(tmpVar_x, tmpVar_y))
+            defs = [MacroDefinition(element_name, element_type)]
+            undefs = [MacroUndef(element_name)]
 
-        return lt_def
+        if isinstance(element_type, (NumpyNDArrayType, HomogeneousTupleType)):
+            defs.append(MacroDefinition(f'{element_name}_rank', element_type.rank))
+            undefs.append(MacroUndef(f'{element_name}_rank'))
+        elif not isinstance(element_type, FixedSizeNumericType):
+            raise NotImplementedError("Support for containers of types defined in other modules is not yet implemented")
+        return defs, undefs
 
     def _build_gFTL_module(self, expr_type):
         """
@@ -583,60 +606,40 @@ def _build_gFTL_module(self, expr_type):
             module = self._generated_gFTL_types[matching_expr_type]
             mod_name = module.name
         else:
-            if isinstance(expr_type, HomogeneousListType):
+            if isinstance(expr_type, (HomogeneousListType, HomogeneousSetType)):
                 element_type = expr_type.element_type
-                if isinstance(element_type, FixedSizeNumericType):
-                    imports_and_macros = [MacroDefinition('T', element_type.primitive_type),
-                              MacroDefinition('T_KINDLEN(context)', KindSpecification(element_type))]
+                if isinstance(expr_type, HomogeneousSetType):
+                    type_name = 'Set'
+                    if not isinstance(element_type, FixedSizeNumericType):
+                        raise NotImplementedError("Support for sets of types which define their own < operator is not yet implemented")
                 else:
-                    imports_and_macros = [MacroDefinition('T', element_type)]
-                if isinstance(element_type, (NumpyNDArrayType, HomogeneousTupleType)):
-                    imports_and_macros.append(MacroDefinition('T_rank', element_type.rank))
-                elif not isinstance(element_type, FixedSizeNumericType):
-                    raise NotImplementedError("Support for lists of types defined in other modules is not yet implemented")
-                imports_and_macros.extend([MacroDefinition('Vector', expr_type),
-                                           MacroDefinition('VectorIterator', IteratorType(expr_type)),
-                                           Import(LiteralString('vector/template.inc'), Module('_', (), ())),
-                                           MacroUndef('Vector'),
-                                           MacroUndef('VectorIterator'),])
-            elif isinstance(expr_type, HomogeneousSetType):
-                element_type = expr_type.element_type
+                    type_name = 'Vector'
                 imports_and_macros = []
-                if isinstance(element_type, FixedSizeNumericType):
-                    lt_def = self._get_comparison_operator(element_type, imports_and_macros)
-                    imports_and_macros.extend([MacroDefinition('T', element_type.primitive_type),
-                              MacroDefinition('T_KINDLEN(context)', KindSpecification(element_type)),
-                              MacroDefinition('T_LT(x,y)', lt_def)])
-                else:
-                    raise NotImplementedError("Support for sets of types which define their own < operator is not yet implemented")
-                imports_and_macros.extend([MacroDefinition('Set', expr_type),
-                                           MacroDefinition('SetIterator', IteratorType(expr_type)),
-                                           Import(LiteralString('set/template.inc'), Module('_', (), ())),
-                                           MacroUndef('Set'),
-                                           MacroUndef('SetIterator')])
+                defs, undefs = self._define_gFTL_element(element_type, imports_and_macros, 'T')
+                imports_and_macros.extend([*defs,
+                                           MacroDefinition(type_name, expr_type),
+                                           MacroDefinition(f'{type_name}Iterator', IteratorType(expr_type)),
+                                           Import(LiteralString(f'{type_name.lower()}/template.inc'), Module('_', (), ())),
+                                           MacroUndef(type_name),
+                                           MacroUndef(f'{type_name}Iterator'),
+                                           *undefs])
             elif isinstance(expr_type, DictType):
                 key_type = expr_type.key_type
                 value_type = expr_type.value_type
                 imports_and_macros = []
-                if isinstance(key_type, FixedSizeNumericType):
-                    lt_def = self._get_comparison_operator(key_type, imports_and_macros)
-                    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:
+                if not isinstance(key_type, FixedSizeNumericType):
                     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.extend([MacroDefinition('Pair', PairType(key_type, value_type)),
+                key_defs, key_undefs = self._define_gFTL_element(key_type, imports_and_macros, 'Key')
+                val_defs, val_undefs = self._define_gFTL_element(value_type, imports_and_macros, 'T')
+                imports_and_macros.extend([*key_defs, *val_defs,
+                                           MacroDefinition('Pair', PairType(key_type, value_type)),
                                            MacroDefinition('Map', expr_type),
                                            MacroDefinition('MapIterator', IteratorType(expr_type)),
                                            Import(LiteralString('map/template.inc'), Module('_', (), ())),
                                            MacroUndef('Pair'),
                                            MacroUndef('Map'),
-                                           MacroUndef('MapIterator')])
+                                           MacroUndef('MapIterator'),
+                                           *key_undefs, *val_undefs])
             else:
                 raise NotImplementedError(f"Unkown gFTL import for type {expr_type}")
 
@@ -3077,6 +3080,19 @@ def _print_PyccelNot(self, expr):
             return '{} == 0'.format(a)
         return '.not. {}'.format(a)
 
+    def _print_PyccelIn(self, expr):
+        container_type = expr.container.class_type
+        element = self._print(expr.element)
+        container = self._print(expr.container)
+        if isinstance(container_type, (HomogeneousSetType, DictType)):
+            return f'{container} % count({element}) /= 0'
+        elif isinstance(container_type, HomogeneousListType):
+            return f'{container} % get_index({element}) /= 0'
+        else:
+            raise errors.report(PYCCEL_RESTRICTION_TODO,
+                    symbol = expr,
+                    severity='fatal')
+
     def _print_Header(self, expr):
         return ''
 

pyccel/codegen/printing/pycode.py

@@ -1091,6 +1091,11 @@ def _print_Duplicate(self, expr):
     def _print_Concatenate(self, expr):
         return ' + '.join([self._print(a) for a in expr.args])
 
+    def _print_PyccelIn(self, expr):
+        element = self._print(expr.element)
+        container = self._print(expr.container)
+        return f'{element} in {container}'
+
     def _print_PyccelSymbol(self, expr):
         return expr
 

pyccel/parser/semantic.py

@@ -98,7 +98,7 @@
 
 from pyccel.ast.operators import PyccelArithmeticOperator, PyccelIs, PyccelIsNot, IfTernaryOperator, PyccelUnarySub
 from pyccel.ast.operators import PyccelNot, PyccelAdd, PyccelMinus, PyccelMul, PyccelPow
-from pyccel.ast.operators import PyccelAssociativeParenthesis, PyccelDiv
+from pyccel.ast.operators import PyccelAssociativeParenthesis, PyccelDiv, PyccelIn
 
 from pyccel.ast.sympy_helper import sympy_to_pyccel, pyccel_to_sympy
 
@@ -2926,6 +2926,25 @@ def _visit_PyccelPow(self, expr):
         else:
             return PyccelPow(base, exponent)
 
+    def _visit_PyccelIn(self, expr):
+        element = self._visit(expr.element)
+        container = self._visit(expr.container)
+        container_type = container.class_type
+        if isinstance(container_type, (DictType, HomogeneousSetType, HomogeneousListType)):
+            element_type = container_type.key_type if isinstance(container_type, DictType) else container_type.element_type
+            if element.class_type == element_type:
+                return PyccelIn(element, container)
+            else:
+                return LiteralFalse()
+
+        container_base = self.scope.find(str(container_type), 'classes') or get_cls_base(container_type)
+        contains_method = container_base.get_method('__contains__', raise_error = isinstance(container_type, CustomDataType))
+        if contains_method:
+            return contains_method(container, element)
+        else:
+            raise errors.report(f"In operator is not yet implemented for type {container_type}",
+                    severity='fatal', symbol=expr)
+
     def _visit_Lambda(self, expr):
         errors.report("Lambda functions are not currently supported",
                 symbol=expr, severity='fatal')

pyccel/parser/syntactic.py

@@ -49,7 +49,7 @@
 from pyccel.ast.operators import PyccelEq,  PyccelNe,  PyccelLt,  PyccelLe,  PyccelGt,  PyccelGe
 from pyccel.ast.operators import PyccelAnd, PyccelOr,  PyccelNot, PyccelMinus
 from pyccel.ast.operators import PyccelUnary, PyccelUnarySub
-from pyccel.ast.operators import PyccelIs, PyccelIsNot
+from pyccel.ast.operators import PyccelIs, PyccelIsNot, PyccelIn
 from pyccel.ast.operators import IfTernaryOperator
 from pyccel.ast.numpyext  import NumpyMatmul
 
@@ -715,6 +715,8 @@ def _visit_Compare(self, stmt):
             return PyccelIs(first, second)
         if isinstance(op, ast.IsNot):
             return PyccelIsNot(first, second)
+        if isinstance(op, ast.In):
+            return PyccelIn(first, second)
 
         return errors.report(PYCCEL_RESTRICTION_UNSUPPORTED_SYNTAX,
                       symbol = stmt,