let numeric operation can find correct function

pyston · Feb 17, 2016 · 989a329 · 989a329
1 parent 746fc43
commit 989a329
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 7 deletions.
diff --git a/src/runtime/float.cpp b/src/runtime/float.cpp
@@ -161,7 +161,7 @@ extern "C" Box* floatDiv(BoxedFloat* lhs, Box* rhs) {
     assert(PyFloat_Check(lhs));
     if (PyInt_Check(rhs)) {
         return floatDivInt(lhs, static_cast<BoxedInt*>(rhs));
-    } else if (PyFloat_Check(rhs)) {
+    } else if (rhs->cls == float_cls) {
         return floatDivFloat(lhs, static_cast<BoxedFloat*>(rhs));
     } else if (PyLong_Check(rhs)) {
         double rhs_f = PyLong_AsDouble(rhs);
@@ -178,7 +178,7 @@ extern "C" Box* floatTruediv(BoxedFloat* lhs, Box* rhs) {
     assert(PyFloat_Check(lhs));
     if (PyInt_Check(rhs)) {
         return floatDivInt(lhs, static_cast<BoxedInt*>(rhs));
-    } else if (PyFloat_Check(rhs)) {
+    } else if (rhs->cls == float_cls) {
         return floatDivFloat(lhs, static_cast<BoxedFloat*>(rhs));
     } else if (PyLong_Check(rhs)) {
         double rhs_f = PyLong_AsDouble(rhs);

diff --git a/src/runtime/int.cpp b/src/runtime/int.cpp
@@ -553,9 +553,9 @@ extern "C" Box* intDiv(BoxedInt* lhs, Box* rhs) {
     if (!PyInt_Check(lhs))
         raiseExcHelper(TypeError, "descriptor '__div__' requires a 'int' object but received a '%s'", getTypeName(lhs));
 
-    if (PyInt_Check(rhs)) {
+    if (rhs->cls == int_cls) {
         return intDivInt(lhs, static_cast<BoxedInt*>(rhs));
-    } else if (PyFloat_Check(rhs)) {
+    } else if (rhs->cls == float_cls) {
         return intDivFloat(lhs, static_cast<BoxedFloat*>(rhs));
     } else {
         return NotImplemented;
@@ -641,9 +641,9 @@ extern "C" Box* intTruediv(BoxedInt* lhs, Box* rhs) {
         raiseExcHelper(TypeError, "descriptor '__truediv__' requires a 'int' object but received a '%s'",
                        getTypeName(lhs));
 
-    if (PyInt_Check(rhs)) {
+    if (rhs->cls == int_cls) {
         return intTruedivInt(lhs, static_cast<BoxedInt*>(rhs));
-    } else if (PyFloat_Check(rhs)) {
+    } else if (rhs->cls == float_cls) {
         return intTruedivFloat(lhs, static_cast<BoxedFloat*>(rhs));
     } else {
         return NotImplemented;

diff --git a/src/runtime/long.cpp b/src/runtime/long.cpp
@@ -1378,7 +1378,7 @@ Box* longTrueDiv(BoxedLong* v1, Box* _v2) {
                        getTypeName(v1));
 
     BoxedLong* v2;
-    if (PyInt_Check(_v2) || PyLong_Check(_v2)) {
+    if (PyInt_Check(_v2) || _v2->cls == long_cls) {
         v2 = (BoxedLong*)PyNumber_Long(_v2);
         if (!v2) {
             throwCAPIException();

diff --git a/src/runtime/objmodel.cpp b/src/runtime/objmodel.cpp
@@ -4565,6 +4565,13 @@ Box* binopInternal(Box* lhs, Box* rhs, int op_type, bool inplace, BinopRewriteAr
                 assert(rewrite_args->out_rtn);
                 rewrite_args->out_success = true;
             }
+
+            if (rhs->cls != lhs->cls && isSubclass(rhs->cls, lhs->cls)) {
+                BoxedString* rop_name = getReverseOpName(op_type);
+                Box* rrtn = callattrInternal1<CXX, REWRITABLE>(rhs, rop_name, CLASS_ONLY, NULL, ArgPassSpec(1), lhs);
+                if (rrtn != NULL && rrtn != NotImplemented)
+                    return rrtn;
+            }
             return lrtn;
         }
     }

diff --git a/test/tests/float.py b/test/tests/float.py
@@ -140,3 +140,35 @@ class F2(float):
                                                                arg2=y)))
             except Exception as e:
                 print(e.message)
+
+
+class Foo1(float):
+    def __rdiv__(self, other):
+        print("float custom operation called")
+        return self / other
+
+
+class Foo2(long):
+    def __rdiv__(self, other):
+        print("long custom operation called")
+        return self / other
+
+
+class Foo3(int):
+    def __rdiv__(self, other):
+        print("int custom operation called")
+        return self / other
+
+a = Foo1(1.5)
+b = Foo2(1L)
+c = Foo3(1)
+
+print(1.5 / a)
+print(1.5 / b)
+print(1.5 / c)
+print(1 / a)
+print(1 / b)
+print(1 / c)
+print(1L / a)
+print(1L / b)
+print(1L / c)