Add support for _len__ (pyccel#2129)

Add support for `__len__` magic method. Fixes pyccel#2106

**Commit Summary**
- Add a `_build_PythonLen` function to handle class variables as
arguments
- Group all magic methods together in `PyClassDef`

Author: EmilyBourne

CHANGELOG.md

@@ -56,6 +56,7 @@ All notable changes to this project will be documented in this file.
 -   #738 : Add support for homogeneous tuples with scalar elements as arguments.
 -   Add a warning about containers in lists.
 -   #2016 : Add support for translating arithmetic magic methods.
+-   #2106 : Add support for `__len__` magic method.
 -   #1980 : Extend The C support for min and max to more than two variables
 -   #2081 : Add support for multi operator expressions
 -   #2061 : Add C support for string declarations.

docs/classes.md

@@ -446,6 +446,7 @@ Pyccel supports a subset of magic methods that are listed here:
 -   `__irshift__`
 -   `__iand__`
 -   `__ior__`
+-   `__len__`
 
 Additionally the following methods are supported in the translation but are lacking the wrapper support that would allow them to be called from Python code:
 

pyccel/ast/builtins.py

@@ -21,7 +21,7 @@
 from .datatypes import HomogeneousTupleType, InhomogeneousTupleType
 from .datatypes import HomogeneousListType, HomogeneousContainerType
 from .datatypes import FixedSizeNumericType, HomogeneousSetType, SymbolicType
-from .datatypes import DictType, VoidType, TypeAlias
+from .datatypes import DictType, VoidType, TypeAlias, StringType
 from .internals import PyccelFunction, Slice, PyccelArrayShapeElement, Iterable
 from .literals  import LiteralInteger, LiteralFloat, LiteralComplex, Nil
 from .literals  import Literal, LiteralImaginaryUnit, convert_to_literal
@@ -725,9 +725,14 @@ class PythonLen(PyccelFunction):
     def __new__(cls, arg):
         if isinstance(arg, LiteralString):
             return LiteralInteger(len(arg.python_value))
+        elif isinstance(arg.class_type, StringType):
+            return super().__new__(cls)
         else:
             return arg.shape[0]
 
+    def __init__(self, arg):
+        super().__init__(arg)
+
 #==============================================================================
 class PythonList(TypedAstNode):
     """

pyccel/ast/cwrapper.py

@@ -718,8 +718,8 @@ class definition.
         wrapped.
     """
     __slots__ = ('_original_class', '_struct_name', '_type_name', '_type_object',
-                 '_new_func', '_properties', '_number_magic_methods')
-    _attribute_nodes = ClassDef._attribute_nodes + ('_number_magic_methods',)
+                 '_new_func', '_properties', '_magic_methods')
+    _attribute_nodes = ClassDef._attribute_nodes + ('_magic_methods',)
 
     def __init__(self, original_class, struct_name, type_name, scope, **kwargs):
         self._original_class = original_class
@@ -728,7 +728,7 @@ def __init__(self, original_class, struct_name, type_name, scope, **kwargs):
         self._type_object = Variable(PyccelPyClassType(), type_name)
         self._new_func = None
         self._properties = ()
-        self._number_magic_methods = ()
+        self._magic_methods = ()
         variables = [Variable(VoidType(), 'instance', memory_handling='alias'),
                      Variable(PyccelPyObject(), 'referenced_objects', memory_handling='alias'),
                      Variable(PythonNativeBool(), 'is_alias')]
@@ -821,12 +821,11 @@ def properties(self):
         """
         return self._properties
 
-    def add_new_magic_number_method(self, method):
+    def add_new_magic_method(self, method):
         """
-        Add a new magic number method to the current class.
+        Add a new magic method to the current class.
 
-        Add a new magic method to the current ClassDef describing
-        a number method.
+        Add a new magic method to the current ClassDef.
 
         Parameters
         ----------
@@ -837,16 +836,16 @@ def add_new_magic_number_method(self, method):
         if not isinstance(method, PyFunctionDef):
             raise TypeError("Method must be FunctionDef")
         method.set_current_user_node(self)
-        self._number_magic_methods += (method,)
+        self._magic_methods += (method,)
 
     @property
-    def number_magic_methods(self):
+    def magic_methods(self):
         """
-        Get the magic methods describing number methods.
+        Get the magic methods describing methods.
 
-        Get the magic methods describing number methods such as __add__.
+        Get the magic methods describing methods such as __add__.
         """
-        return self._number_magic_methods
+        return self._magic_methods
 
 #-------------------------------------------------------------------
 

pyccel/ast/itertoolsext.py

@@ -5,7 +5,6 @@
 """
 This module represent a call to the itertools functions for code generation.
 """
-from .builtins  import PythonLen, PythonRange
 from .core      import PyccelFunctionDef, Module
 from .internals import Iterable
 

pyccel/codegen/printing/cwrappercode.py

@@ -286,7 +286,7 @@ def _print_ModuleHeader(self, expr):
             sig_methods = c.methods + (c.new_func,) + tuple(f for i in c.interfaces for f in i.functions) + \
                           tuple(i.interface_func for i in c.interfaces) + \
                           tuple(getset for p in c.properties for getset in (p.getter, p.setter) if getset) + \
-                          c.number_magic_methods
+                          c.magic_methods
             function_signatures += '\n'+''.join(self.function_signature(f)+';\n' for f in sig_methods)
             macro_defs += f'#define {type_name} (*(PyTypeObject*){API_var.name}[{i}])\n'
 
@@ -396,7 +396,7 @@ def _print_PyClassDef(self, expr):
         getters = tuple(p.getter for p in expr.properties)
         setters = tuple(p.setter for p in expr.properties if p.setter)
         print_methods = expr.methods + (expr.new_func,) + expr.interfaces + \
-                         expr.number_magic_methods + getters + setters
+                         expr.magic_methods + getters + setters
         functions = '\n'.join(self._print(f) for f in print_methods)
         init_string = ''
         del_string = ''
@@ -435,44 +435,58 @@ def _print_PyClassDef(self, expr):
                                      '},\n')
                                      for name, (wrapper_name, doc_string) in funcs.items())
 
+        magic_methods = {self.get_python_name(original_scope, f.original_function): f for f in expr.magic_methods}
+
         number_magic_method_name = self.scope.get_new_name(f'{expr.name}_number_methods')
-        number_magic_methods = {self.get_python_name(original_scope, f.original_function): f for f in expr.number_magic_methods}
 
         number_magic_methods_def = f"static PyNumberMethods {number_magic_method_name} = {{\n"
-        if '__add__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_add = (binaryfunc){number_magic_methods['__add__'].name},\n"
-        if '__sub__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_subtract = (binaryfunc){number_magic_methods['__sub__'].name},\n"
-        if '__mul__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_multiply = (binaryfunc){number_magic_methods['__mul__'].name},\n"
-        if '__truediv__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_true_divide = (binaryfunc){number_magic_methods['__truediv__'].name},\n"
-        if '__lshift__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_lshift = (binaryfunc){number_magic_methods['__lshift__'].name},\n"
-        if '__rshift__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_rshift = (binaryfunc){number_magic_methods['__rshift__'].name},\n"
-        if '__and__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_and = (binaryfunc){number_magic_methods['__and__'].name},\n"
-        if '__or__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_or = (binaryfunc){number_magic_methods['__or__'].name},\n"
-        if '__iadd__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_inplace_add = (binaryfunc){number_magic_methods['__iadd__'].name},\n"
-        if '__isub__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_inplace_subtract = (binaryfunc){number_magic_methods['__isub__'].name},\n"
-        if '__imul__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_inplace_multiply = (binaryfunc){number_magic_methods['__imul__'].name},\n"
-        if '__itruediv__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_inplace_true_divide = (binaryfunc){number_magic_methods['__itruediv__'].name},\n"
-        if '__ilshift__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_inplace_lshift = (binaryfunc){number_magic_methods['__ilshift__'].name},\n"
-        if '__irshift__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_inplace_rshift = (binaryfunc){number_magic_methods['__irshift__'].name},\n"
-        if '__iand__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_inplace_and = (binaryfunc){number_magic_methods['__iand__'].name},\n"
-        if '__ior__' in number_magic_methods:
-            number_magic_methods_def += f"     .nb_inplace_or = (binaryfunc){number_magic_methods['__ior__'].name},\n"
+        if '__add__' in magic_methods:
+            number_magic_methods_def += f"     .nb_add = (binaryfunc){magic_methods['__add__'].name},\n"
+        if '__sub__' in magic_methods:
+            number_magic_methods_def += f"     .nb_subtract = (binaryfunc){magic_methods['__sub__'].name},\n"
+        if '__mul__' in magic_methods:
+            number_magic_methods_def += f"     .nb_multiply = (binaryfunc){magic_methods['__mul__'].name},\n"
+        if '__truediv__' in magic_methods:
+            number_magic_methods_def += f"     .nb_true_divide = (binaryfunc){magic_methods['__truediv__'].name},\n"
+        if '__lshift__' in magic_methods:
+            number_magic_methods_def += f"     .nb_lshift = (binaryfunc){magic_methods['__lshift__'].name},\n"
+        if '__rshift__' in magic_methods:
+            number_magic_methods_def += f"     .nb_rshift = (binaryfunc){magic_methods['__rshift__'].name},\n"
+        if '__and__' in magic_methods:
+            number_magic_methods_def += f"     .nb_and = (binaryfunc){magic_methods['__and__'].name},\n"
+        if '__or__' in magic_methods:
+            number_magic_methods_def += f"     .nb_or = (binaryfunc){magic_methods['__or__'].name},\n"
+        if '__iadd__' in magic_methods:
+            number_magic_methods_def += f"     .nb_inplace_add = (binaryfunc){magic_methods['__iadd__'].name},\n"
+        if '__isub__' in magic_methods:
+            number_magic_methods_def += f"     .nb_inplace_subtract = (binaryfunc){magic_methods['__isub__'].name},\n"
+        if '__imul__' in magic_methods:
+            number_magic_methods_def += f"     .nb_inplace_multiply = (binaryfunc){magic_methods['__imul__'].name},\n"
+        if '__itruediv__' in magic_methods:
+            number_magic_methods_def += f"     .nb_inplace_true_divide = (binaryfunc){magic_methods['__itruediv__'].name},\n"
+        if '__ilshift__' in magic_methods:
+            number_magic_methods_def += f"     .nb_inplace_lshift = (binaryfunc){magic_methods['__ilshift__'].name},\n"
+        if '__irshift__' in magic_methods:
+            number_magic_methods_def += f"     .nb_inplace_rshift = (binaryfunc){magic_methods['__irshift__'].name},\n"
+        if '__iand__' in magic_methods:
+            number_magic_methods_def += f"     .nb_inplace_and = (binaryfunc){magic_methods['__iand__'].name},\n"
+        if '__ior__' in magic_methods:
+            number_magic_methods_def += f"     .nb_inplace_or = (binaryfunc){magic_methods['__ior__'].name},\n"
         number_magic_methods_def += '};\n'
 
+        seq_magic_method_name = self.scope.get_new_name(f'{expr.name}_sequence_methods')
+
+        seq_magic_methods_def = f"static PySequenceMethods {seq_magic_method_name} = {{\n"
+        if '__len__' in magic_methods:
+            seq_magic_methods_def += f"    .sq_length = {magic_methods['__len__'].name},\n"
+        seq_magic_methods_def += '};\n'
+
+        map_magic_method_name = self.scope.get_new_name(f'{expr.name}_mapping_methods')
+        map_magic_methods_def = f"static PyMappingMethods {map_magic_method_name} = {{\n"
+        if '__len__' in magic_methods:
+            map_magic_methods_def += f"    .mp_length = {magic_methods['__len__'].name},\n"
+        map_magic_methods_def += '};\n'
+
         method_def_name = self.scope.get_new_name(f'{expr.name}_methods')
         method_def = (f'static PyMethodDef {method_def_name}[] = {{\n'
                         f'{method_def_funcs}'
@@ -488,6 +502,8 @@ def _print_PyClassDef(self, expr):
                 "    PyVarObject_HEAD_INIT(NULL, 0)\n"
                 f"    .tp_name = \"{self._module_name}.{name}\",\n"
                 f"    .tp_as_number = &{number_magic_method_name},\n"
+                f"    .tp_as_sequence = &{seq_magic_method_name},\n"
+                f"    .tp_as_mapping = &{map_magic_method_name},\n"
                 f"    .tp_doc = PyDoc_STR({docstring}),\n"
                 f"    .tp_basicsize = sizeof(struct {struct_name}),\n"
                  "    .tp_itemsize = 0,\n"
@@ -498,7 +514,8 @@ def _print_PyClassDef(self, expr):
                 f"    .tp_getset = {property_def_name},\n"
                 "};\n")
 
-        return '\n'.join((method_def, number_magic_methods_def, property_def, type_code, functions))
+        return '\n'.join((method_def, number_magic_methods_def, seq_magic_methods_def,
+                        map_magic_methods_def, property_def, type_code, functions))
 
     def _print_PyModInitFunc(self, expr):
         decs = ''.join(self._print(d) for d in expr.declarations)

pyccel/codegen/wrapper/c_to_python_wrapper.py

@@ -1498,7 +1498,7 @@ def _wrap_FunctionDef(self, expr):
             func_args = [FunctionDefArgument(a) for a in func_args]
             body = []
         else:
-            if in_interface or original_func_name in magic_binary_funcs:
+            if in_interface or original_func_name in magic_binary_funcs or original_func_name == '__len__':
                 func_args = [FunctionDefArgument(a) for a in self._get_python_argument_variables(python_args)]
                 body = []
             else:
@@ -1549,7 +1549,12 @@ def _wrap_FunctionDef(self, expr):
                     body.append(If( IfSection(PyccelIsNot(v, Nil()), [Deallocate(v)]) ))
                 else:
                     body.append(Deallocate(v))
-        body.extend(wrapped_results['body'])
+
+        if original_func_name == '__len__':
+            self.scope.remove_variable(python_result_variable)
+            python_result_variable = c_results[0]
+        else:
+            body.extend(wrapped_results['body'])
         body.extend(ai for arg in wrapped_args for ai in arg['clean_up'])
 
         # Pack the Python compatible results of the function into one argument.
@@ -2040,8 +2045,8 @@ def _wrap_ClassDef(self, expr):
                 wrapped_class.add_new_method(self._get_class_destructor(f, orig_cls_dtype, wrapped_class.scope))
             elif python_name == '__init__':
                 wrapped_class.add_new_method(self._get_class_initialiser(f, orig_cls_dtype))
-            elif python_name in magic_binary_funcs:
-                wrapped_class.add_new_magic_number_method(self._wrap(f))
+            elif python_name in (*magic_binary_funcs, '__len__'):
+                wrapped_class.add_new_magic_method(self._wrap(f))
             elif 'property' in f.decorators:
                 wrapped_class.add_property(self._wrap(f))
             else:

pyccel/parser/semantic.py

@@ -30,8 +30,8 @@
 from pyccel.ast.builtins import PythonComplex, PythonDict, PythonDictFunction, PythonListFunction
 from pyccel.ast.builtins import builtin_functions_dict, PythonImag, PythonReal
 from pyccel.ast.builtins import PythonList, PythonConjugate , PythonSet, VariableIterator
-from pyccel.ast.builtins import (PythonRange, PythonZip, PythonEnumerate,
-                                 PythonTuple, Lambda, PythonMap)
+from pyccel.ast.builtins import PythonRange, PythonZip, PythonEnumerate, PythonTuple
+from pyccel.ast.builtins import Lambda, PythonMap
 
 from pyccel.ast.builtin_methods.list_methods import ListMethod, ListAppend
 from pyccel.ast.builtin_methods.set_methods  import SetAdd, SetUnion, SetCopy, SetIntersectionUpdate
@@ -5508,3 +5508,43 @@ def _build_SetIntersection(self, expr, function_call_args):
             self._additional_exprs[-1].extend(body)
             return lhs
 
+    def _build_PythonLen(self, expr, function_call_args):
+        """
+        Method to visit a PythonLen node.
+
+        The purpose of this `_build` method is to construct a node representing
+        a call to the PythonLen function. This function returns the first element
+        of the shape of a variable, or a call to a method which calculates the
+        length (e.g. the `__len__` function).
+
+        Parameters
+        ----------
+        expr : DottedName
+            The syntactic node that represent the call to `len()`.
+
+        function_call_args : iterable[FunctionCallArgument]
+            The semantic arguments passed to the function.
+
+        Returns
+        -------
+        TypedAstNode
+            The node representing an object which allows the result of the
+            PythonLen function to be obtained.
+        """
+        arg = function_call_args[0].value
+        class_type = arg.class_type
+        if isinstance(arg, LiteralString):
+            return LiteralInteger(len(arg.python_value))
+        elif isinstance(arg.class_type, CustomDataType):
+            class_base = self.scope.find(str(class_type), 'classes') or get_cls_base(class_type)
+            magic_method = class_base.get_method('__len__', False)
+            if magic_method:
+                return self._handle_function(expr, magic_method, function_call_args)
+            else:
+                raise errors.report(f"__len__ not implemented for type {class_type}",
+                        severity='fatal', symbol=expr)
+        elif arg.rank > 0:
+            return arg.shape[0]
+        else:
+            raise errors.report(f"__len__ not implemented for type {class_type}",
+                    severity='fatal', symbol=expr)

tests/epyccel/classes/class_magic.py

@@ -59,3 +59,6 @@ def __iand__(self, other : int):
     def __ior__(self, other : int):
         self.x |= other
         return self
+
+    def __len__(self):
+        return 1

tests/epyccel/test_epyccel_classes.py

@@ -444,3 +444,4 @@ def test_class_magic(language):
 
     assert a_py.x == a_l.x
 
+    assert len(a_py) == len(a_l)