Add wrapping support for arithmetic magic methods (pyccel#2128)

Add wrapping support for arithmetic magic methods to allow them to be
called from Python. Fixes pyccel#2024

Author: EmilyBourne

CHANGELOG.md

@@ -54,7 +54,7 @@ All notable changes to this project will be documented in this file.
 -   #1583 : Allow inhomogeneous tuples in classes.
 -   #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 (methods cannot yet be used from Python).
+-   #2016 : Add support for translating arithmetic magic methods.
 -   #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

@@ -2,6 +2,15 @@
 
 Pyccel strives to provide robust support for object-oriented programming concepts commonly used by developers. In Pyccel, classes are a fundamental building block for creating structured and reusable code. This documentation outlines key features and considerations when working with classes in Pyccel.
 
+## Contents
+
+1. [Constructor Method](#constructor-method)
+2. [Destructor Method](#destructor-method)
+3. [Class Methods](#class-methods)
+4. [Class Properties](#class-properties)
+5. [Magic Methods](#magic-methods)
+6. [Limitations](#limitations)
+
 ## Constructor Method
 
 -   The Constructor Method, `__init__`, is used to initialise the object's attributes.
@@ -188,7 +197,7 @@ MyClass Object created!
 ==158858== ERROR SUMMARY: 0 errors from 0 contexts (suppressed: 0 from 0)
 ```
 
-## Class properties
+## Class Properties
 
 Pyccel now supports class properties (to retrieve a constant value only).
 
@@ -415,6 +424,42 @@ int main()
   call obj % free()
 ```
 
+## Magic Methods
+
+Pyccel supports a subset of magic methods that are listed here:
+
+-   `__add__`
+-   `__sub__`
+-   `__mul__`
+-   `__truediv__`
+-   `__pow__`
+-   `__lshift__`
+-   `__rshift__`
+-   `__and__`
+-   `__or__`
+-   `__iadd__`
+-   `__isub__`
+-   `__imul__`
+-   `__itruediv__`
+-   `__ipow__`
+-   `__ilshift__`
+-   `__irshift__`
+-   `__iand__`
+-   `__ior__`
+
+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:
+
+-   `__radd__`
+-   `__rsub__`
+-   `__rmul__`
+-   `__rtruediv__`
+-   `__rpow__`
+-   `__rlshift__`
+-   `__rrshift__`
+-   `__rand__`
+-   `__ror__`
+-   `__contains__`
+
 ## Limitations
 
-It's important to note that Pyccel does not support class inheritance, magic methods or static class variables. For our first implementation, the focus of Pyccel is primarily on core class functionality and memory management.
+It's important to note that Pyccel does not support class inheritance, or static class variables. For our first implementation, the focus of Pyccel is primarily on core class functionality and memory management.

pyccel/ast/core.py

@@ -13,7 +13,7 @@
 
 from .basic     import PyccelAstNode, TypedAstNode, iterable, ScopedAstNode
 
-from .bitwise_operators import PyccelBitOr, PyccelBitAnd
+from .bitwise_operators import PyccelBitOr, PyccelBitAnd, PyccelLShift, PyccelRShift
 
 from .builtins  import PythonBool, PythonTuple
 
@@ -750,6 +750,8 @@ class AugAssign(Assign):
             '%' : PyccelMod,
             '|' : PyccelBitOr,
             '&' : PyccelBitAnd,
+            '<<': PyccelLShift,
+            '>>': PyccelRShift,
         }
 
     def __init__(

pyccel/ast/cwrapper.py

@@ -718,7 +718,8 @@ class definition.
         wrapped.
     """
     __slots__ = ('_original_class', '_struct_name', '_type_name', '_type_object',
-                 '_new_func', '_properties')
+                 '_new_func', '_properties', '_number_magic_methods')
+    _attribute_nodes = ClassDef._attribute_nodes + ('_number_magic_methods',)
 
     def __init__(self, original_class, struct_name, type_name, scope, **kwargs):
         self._original_class = original_class
@@ -727,6 +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 = ()
         variables = [Variable(VoidType(), 'instance', memory_handling='alias'),
                      Variable(PyccelPyObject(), 'referenced_objects', memory_handling='alias'),
                      Variable(PythonNativeBool(), 'is_alias')]
@@ -819,6 +821,33 @@ def properties(self):
         """
         return self._properties
 
+    def add_new_magic_number_method(self, method):
+        """
+        Add a new magic number method to the current class.
+
+        Add a new magic method to the current ClassDef describing
+        a number method.
+
+        Parameters
+        ----------
+        method : FunctionDef
+            The Method that will be added.
+        """
+
+        if not isinstance(method, PyFunctionDef):
+            raise TypeError("Method must be FunctionDef")
+        method.set_current_user_node(self)
+        self._number_magic_methods += (method,)
+
+    @property
+    def number_magic_methods(self):
+        """
+        Get the magic methods describing number methods.
+
+        Get the magic methods describing number methods such as __add__.
+        """
+        return self._number_magic_methods
+
 #-------------------------------------------------------------------
 
 class PyGetSetDefElement(PyccelAstNode):

pyccel/codegen/printing/cwrappercode.py

@@ -285,7 +285,8 @@ def _print_ModuleHeader(self, expr):
                     "};\n")
             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)
+                          tuple(getset for p in c.properties for getset in (p.getter, p.setter) if getset) + \
+                          c.number_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'
 
@@ -395,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 + \
-                         getters + setters
+                         expr.number_magic_methods + getters + setters
         functions = '\n'.join(self._print(f) for f in print_methods)
         init_string = ''
         del_string = ''
@@ -434,6 +435,44 @@ def _print_PyClassDef(self, expr):
                                      '},\n')
                                      for name, (wrapper_name, doc_string) in funcs.items())
 
+        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"
+        number_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}'
@@ -448,6 +487,7 @@ def _print_PyClassDef(self, expr):
         type_code = (f"static PyTypeObject {type_name} = {{\n"
                 "    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_doc = PyDoc_STR({docstring}),\n"
                 f"    .tp_basicsize = sizeof(struct {struct_name}),\n"
                  "    .tp_itemsize = 0,\n"
@@ -458,7 +498,7 @@ def _print_PyClassDef(self, expr):
                 f"    .tp_getset = {property_def_name},\n"
                 "};\n")
 
-        return '\n'.join((method_def, property_def, type_code, functions))
+        return '\n'.join((method_def, number_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/printing/pycode.py

@@ -21,6 +21,8 @@
 from pyccel.ast.utilities  import builtin_import_registry as pyccel_builtin_import_registry
 from pyccel.ast.utilities  import decorators_mod
 
+from pyccel.parser.semantic import magic_method_map
+
 from pyccel.codegen.printing.codeprinter import CodePrinter
 
 from pyccel.errors.errors import Errors
@@ -299,11 +301,12 @@ def _print_FunctionDef(self, expr):
 
         body = ''.join([docstring, functions, interfaces, imports, body])
 
-        code = ('def {name}({args}):\n'
-                '{body}\n').format(
-                        name=name,
-                        args=args,
-                        body=body)
+        # Put back return removed in semantic stage
+        if name.startswith('__i') and ('__'+name[3:]) in magic_method_map.values():
+            body += f'    return {expr.arguments[0].name}\n'
+
+        code = (f'def {name}({args}):\n'
+                f'{body}\n')
         decorators = expr.decorators.copy()
         if decorators:
             if decorators['template']:

pyccel/codegen/wrapper/c_to_python_wrapper.py

@@ -64,6 +64,26 @@
 cwrapper_ndarray_imports = [Import('cwrapper_ndarrays', Module('cwrapper_ndarrays', (), ())),
                             Import('ndarrays', Module('ndarrays', (), ()))]
 
+magic_binary_funcs = ('__add__',
+                      '__sub__',
+                      '__mul__',
+                      '__truediv__',
+                      '__pow__',
+                      '__lshift__',
+                      '__rshift__',
+                      '__and__',
+                      '__or__',
+                      '__iadd__',
+                      '__isub__',
+                      '__imul__',
+                      '__itruediv__',
+                      '__ipow__',
+                      '__ilshift__',
+                      '__irshift__',
+                      '__iand__',
+                      '__ior__',
+                      )
+
 class CToPythonWrapper(Wrapper):
     """
     Class for creating a wrapper exposing C code to Python.
@@ -1409,7 +1429,7 @@ def _wrap_Interface(self, expr):
 
     def _wrap_FunctionDef(self, expr):
         """
-        Build a `PyFunctionDef` form a `FunctionDef`.
+        Build a `PyFunctionDef` from a `FunctionDef`.
 
         Create a `PyFunctionDef` which wraps a C-compatible `FunctionDef`.
         The `PyFunctionDef` should take three arguments (`self`, `args`,
@@ -1431,6 +1451,7 @@ def _wrap_FunctionDef(self, expr):
         func_name = self.scope.get_new_name(expr.name+'_wrapper')
         func_scope = self.scope.new_child_scope(func_name)
         self.scope = func_scope
+        original_func_name = original_func.scope.get_python_name(original_func.name)
 
         possible_class_base = expr.get_user_nodes((ClassDef,))
         if possible_class_base:
@@ -1477,7 +1498,7 @@ def _wrap_FunctionDef(self, expr):
             func_args = [FunctionDefArgument(a) for a in func_args]
             body = []
         else:
-            if in_interface:
+            if in_interface or original_func_name in magic_binary_funcs:
                 func_args = [FunctionDefArgument(a) for a in self._get_python_argument_variables(python_args)]
                 body = []
             else:
@@ -1490,7 +1511,12 @@ def _wrap_FunctionDef(self, expr):
 
         # Get the code required to wrap the C-compatible results into Python objects
         # This function creates variables so it must be called before extracting them from the scope.
-        if len(python_results) == 0:
+        if original_func_name in magic_binary_funcs and original_func_name.startswith('__i'):
+            res = func_args[0].var.clone(self.scope.get_new_name(func_args[0].var.name), is_argument=False)
+            wrapped_results = {'c_results': [], 'py_result': res, 'body': []}
+            body.append(AliasAssign(res, func_args[0].var))
+            body.append(Py_INCREF(res))
+        elif len(python_results) == 0:
             wrapped_results = {'c_results': [], 'py_result': Py_None, 'body': []}
         elif len(python_results) == 1:
             wrapped_results = self._extract_FunctionDefResult(original_func.results[0].var, is_bind_c_function_def, expr)
@@ -2014,6 +2040,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 'property' in f.decorators:
                 wrapped_class.add_property(self._wrap(f))
             else:

pyccel/parser/syntactic.py

@@ -444,6 +444,10 @@ def _visit_AugAssign(self, stmt):
             return AugAssign(lhs, '|', rhs)
         elif isinstance(stmt.op, ast.BitAnd):
             return AugAssign(lhs, '&', rhs)
+        elif isinstance(stmt.op, ast.LShift):
+            return AugAssign(lhs, '<<', rhs)
+        elif isinstance(stmt.op, ast.RShift):
+            return AugAssign(lhs, '>>', rhs)
         else:
             return errors.report(PYCCEL_RESTRICTION_TODO, symbol = stmt,
                     severity='error')