Add support for pow operator.

sairam4123 · May 24, 2020 · d849fb2 · d849fb2
1 parent e6a8124
commit d849fb2
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 3 deletions.
diff --git a/Cython/Compiler/ModuleNode.py b/Cython/Compiler/ModuleNode.py
@@ -1904,6 +1904,13 @@ def generate_binop_function(self, scope, slot, code):
         if scope.directives['c_api_binop_methods']:
             code.putln('#define %s %s' % (func_name, slot.left_slot.slot_code(scope)))
         else:
+            if slot.left_slot.signature == TypeSlots.binaryfunc:
+                extra_arg = extra_arg_decl = ''
+            elif slot.left_slot.signature == TypeSlots.ternaryfunc:
+                extra_arg = ', extra_arg'
+                extra_arg_decl = ', PyObject* extra_arg'
+            else:
+                error(entry.pos, "Unexpected type lost signature: %s" % slot)
             def has_slot_method(method_name):
                 entry = scope.lookup(method_name)
                 return bool(entry and entry.is_special and entry.func_cname)
@@ -1914,7 +1921,7 @@ def call_slot_method(method_name, reverse):
                 else:
                     operands = "left, right"
                 if entry and entry.is_special and entry.func_cname:
-                    return "%s(%s)" % (entry.func_cname, operands)
+                    return "%s(%s%s)" % (entry.func_cname, operands, extra_arg)
                 else:
                     py_ident = code.intern_identifier(EncodedString(method_name))
                     return "%s_maybe_call_super(%s, %s)" % (func_name, operands, py_ident)
@@ -1928,6 +1935,8 @@ def call_slot_method(method_name, reverse):
                         "call_left": call_slot_method(slot.left_slot.method_name, reverse=False),
                         "call_right": call_slot_method(slot.right_slot.method_name, reverse=True),
                         "type_cname": '((PyTypeObject*) %s)' % scope.namespace_cname,
+                        "extra_arg": extra_arg,
+                        "extra_arg_decl": extra_arg_decl,
                         }).impl.strip())
             code.putln()
         if preprocessor_guard:

diff --git a/Cython/Utility/ExtensionTypes.c b/Cython/Utility/ExtensionTypes.c
@@ -281,7 +281,7 @@ static int __Pyx_setup_reduce(PyObject* type_obj) {
 
 /////////////// BinopSlot ///////////////
 
-static CYTHON_INLINE PyObject *{{func_name}}_maybe_call_super(PyObject *self, PyObject *other, PyObject* name) {
+static CYTHON_INLINE PyObject *{{func_name}}_maybe_call_super(PyObject *self, PyObject *other, PyObject* name {{extra_arg_decl}}) {
     PyObject *res;
     PyObject *method;
     if (!Py_TYPE(self)->tp_base) {
@@ -293,15 +293,19 @@ static CYTHON_INLINE PyObject *{{func_name}}_maybe_call_super(PyObject *self, Py
         PyErr_Clear();
         return Py_INCREF(Py_NotImplemented), Py_NotImplemented;
     }
+    #if {{int(not extra_arg)}}
     res = __Pyx_PyObject_Call2Args(method, self, other);
+    #else
+    res = PyObject_CallFunctionObjArgs(method, self, other {{extra_arg}}, NULL);
+    #endif
     Py_DECREF(method);
     if (!res) {
         return Py_INCREF(Py_NotImplemented), Py_NotImplemented;
     }
     return res;
 }
 
-static PyObject *{{func_name}}(PyObject *left, PyObject *right) {
+static PyObject *{{func_name}}(PyObject *left, PyObject *right {{extra_arg_decl}}) {
     PyObject *res;
     int maybe_self_is_left, maybe_self_is_right = 0;
     maybe_self_is_left = Py_TYPE(left) == Py_TYPE(right)

diff --git a/tests/run/binop_reverse_methods_GH2056.pyx b/tests/run/binop_reverse_methods_GH2056.pyx
@@ -8,13 +8,26 @@ class Base(object):
     'Base.__add__(Base(), 2)'
     >>> 2 + Base()
     'Base.__radd__(Base(), 2)'
+
+    >>> Base() ** 2
+    'Base.__pow__(Base(), 2, None)'
+    >>> 2 ** Base()
+    'Base.__rpow__(Base(), 2, None)'
+    >>> pow(Base(), 2, 100)
+    'Base.__pow__(Base(), 2, 100)'
     """
     def __add__(self, other):
         return "Base.__add__(%s, %s)" % (self, other)
 
     def __radd__(self, other):
         return "Base.__radd__(%s, %s)" % (self, other)
 
+    def __pow__(self, other, mod):
+        return "Base.__pow__(%s, %s, %s)" % (self, other, mod)
+
+    def __rpow__(self, other, mod):
+        return "Base.__rpow__(%s, %s, %s)" % (self, other, mod)
+
     def __repr__(self):
         return "%s()" % (self.__class__.__name__)