py/objtype: Handle NotImplemented return from binary special methods.

NotImplemented means "try other fallbacks (like calling __rop__
instead of __op__) and if nothing works, raise TypeError". As
MicroPython doesn't implement any fallbacks, signal to raise
TypeError right away.

3 files changed, 67 insertions, 2 deletions

diff --git a/py/obj.h b/py/obj.h
index 2adcab1a6..323423b3e 100644
--- a/py/obj.h
+++ b/py/obj.h
@@ -616,6 +616,7 @@ extern const mp_obj_type_t mp_type_ZeroDivisionError;
 #define mp_const_true (MP_OBJ_FROM_PTR(&mp_const_true_obj))
 #define mp_const_empty_bytes (MP_OBJ_FROM_PTR(&mp_const_empty_bytes_obj))
 #define mp_const_empty_tuple (MP_OBJ_FROM_PTR(&mp_const_empty_tuple_obj))
+#define mp_const_notimplemented (MP_OBJ_FROM_PTR(&mp_const_notimplemented_obj))
 extern const struct _mp_obj_none_t mp_const_none_obj;
 extern const struct _mp_obj_bool_t mp_const_false_obj;
 extern const struct _mp_obj_bool_t mp_const_true_obj;
diff --git a/py/objtype.c b/py/objtype.c
index 08e276d72..4973236a4 100644
--- a/py/objtype.c
+++ b/py/objtype.c
@@ -471,14 +471,28 @@ STATIC mp_obj_t instance_binary_op(mp_binary_op_t op, mp_obj_t lhs_in, mp_obj_t
         .is_type = false,
     };
     mp_obj_class_lookup(&lookup, lhs->base.type);
+
+    mp_obj_t res;
     if (dest[0] == MP_OBJ_SENTINEL) {
-        return mp_binary_op(op, lhs->subobj[0], rhs_in);
+        res = mp_binary_op(op, lhs->subobj[0], rhs_in);
     } else if (dest[0] != MP_OBJ_NULL) {
         dest[2] = rhs_in;
-        return mp_call_method_n_kw(1, 0, dest);
+        res = mp_call_method_n_kw(1, 0, dest);
     } else {
         return MP_OBJ_NULL; // op not supported
     }
+
+    #if MICROPY_PY_BUILTINS_NOTIMPLEMENTED
+    // NotImplemented means "try other fallbacks (like calling __rop__
+    // instead of __op__) and if nothing works, raise TypeError". As
+    // MicroPython doesn't implement any fallbacks, signal to raise
+    // TypeError right away.
+    if (res == mp_const_notimplemented) {
+        return MP_OBJ_NULL; // op not supported
+    }
+    #endif
+
+    return res;
 }
 
 STATIC void mp_obj_instance_load_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
diff --git a/tests/basics/class_notimpl.py b/tests/basics/class_notimpl.py
new file mode 100644
index 000000000..7fd8166f9
--- /dev/null
+++ b/tests/basics/class_notimpl.py
@@ -0,0 +1,50 @@
+# Test that returning of NotImplemented from binary op methods leads to
+# TypeError.
+try:
+    NotImplemented
+except NameError:
+    print("SKIP")
+    raise SystemExit
+
+class C:
+    def __init__(self, value):
+        self.value = value
+
+    def __str__(self):
+        return "C(%s)" % self.value
+
+    def __add__(self, rhs):
+        print(self, '+', rhs)
+        return NotImplemented
+
+    def __sub__(self, rhs):
+        print(self, '-', rhs)
+        return NotImplemented
+
+    def __lt__(self, rhs):
+        print(self, '<', rhs)
+        return NotImplemented
+
+    def __neg__(self):
+        print('-', self)
+        return NotImplemented
+
+c = C(0)
+
+try:
+    c + 1
+except TypeError:
+    print("TypeError")
+
+try:
+    c - 2
+except TypeError:
+    print("TypeError")
+
+try:
+    c < 1
+except TypeError:
+    print("TypeError")
+
+# NotImplemented isn't handled specially in unary methods
+print(-c)