py/obj: Add support for __float__ and __complex__ functions.

pi-anl · dpgeorge · commit 1e87b56219c6 · 2022-07-25T14:23:34.000+10:00
diff --git a/py/obj.c b/py/obj.c
@@ -355,9 +355,13 @@ bool mp_obj_get_float_maybe(mp_obj_t arg, mp_float_t *value) {
     } else if (mp_obj_is_float(arg)) {
         val = mp_obj_float_get(arg);
     } else {
-        return false;
+        arg = mp_unary_op(MP_UNARY_OP_FLOAT_MAYBE, (mp_obj_t)arg);
+        if (arg != MP_OBJ_NULL && mp_obj_is_float(arg)) {
+            val = mp_obj_float_get(arg);
+        } else {
+            return false;
+        }
     }
-
     *value = val;
     return true;
 }
@@ -399,7 +403,12 @@ bool mp_obj_get_complex_maybe(mp_obj_t arg, mp_float_t *real, mp_float_t *imag)
     } else if (mp_obj_is_type(arg, &mp_type_complex)) {
         mp_obj_complex_get(arg, real, imag);
     } else {
-        return false;
+        arg = mp_unary_op(MP_UNARY_OP_COMPLEX_MAYBE, (mp_obj_t)arg);
+        if (arg != MP_OBJ_NULL && mp_obj_is_type(arg, &mp_type_complex)) {
+            mp_obj_complex_get(arg, real, imag);
+        } else {
+            return false;
+        }
     }
     return true;
 }
diff --git a/py/objcomplex.c b/py/objcomplex.c
@@ -88,6 +88,10 @@ STATIC mp_obj_t complex_make_new(const mp_obj_type_t *type_in, size_t n_args, si
                 // a complex, just return it
                 return args[0];
             } else {
+                mp_float_t real, imag;
+                if (mp_obj_get_complex_maybe(args[0], &real, &imag)) {
+                    return mp_obj_new_complex(real, imag);
+                }
                 // something else, try to cast it to a complex
                 return mp_obj_new_complex(mp_obj_get_float(args[0]), 0);
             }
diff --git a/py/objtype.c b/py/objtype.c
@@ -378,6 +378,12 @@ const byte mp_unary_op_method_name[MP_UNARY_OP_NUM_RUNTIME] = {
     [MP_UNARY_OP_INVERT] = MP_QSTR___invert__,
     [MP_UNARY_OP_ABS] = MP_QSTR___abs__,
     #endif
+    #if MICROPY_PY_BUILTINS_FLOAT
+    [MP_UNARY_OP_FLOAT_MAYBE] = MP_QSTR___float__,
+    #if MICROPY_PY_BUILTINS_COMPLEX
+    [MP_UNARY_OP_COMPLEX_MAYBE] = MP_QSTR___complex__,
+    #endif
+    #endif
     #if MICROPY_PY_SYS_GETSIZEOF
     [MP_UNARY_OP_SIZEOF] = MP_QSTR___sizeof__,
     #endif
diff --git a/py/runtime.c b/py/runtime.c
@@ -319,6 +319,15 @@ mp_obj_t mp_unary_op(mp_unary_op_t op, mp_obj_t arg) {
             // if arg==mp_const_none.
             return mp_const_true;
         }
+        #if MICROPY_PY_BUILTINS_FLOAT
+        if (op == MP_UNARY_OP_FLOAT_MAYBE
+            #if MICROPY_PY_BUILTINS_COMPLEX
+            || op == MP_UNARY_OP_COMPLEX_MAYBE
+            #endif
+            ) {
+            return MP_OBJ_NULL;
+        }
+        #endif
         // With MP_UNARY_OP_INT, mp_unary_op() becomes a fallback for mp_obj_get_int().
         // In this case provide a more focused error message to not confuse, e.g. chr(1.0)
         #if MICROPY_ERROR_REPORTING <= MICROPY_ERROR_REPORTING_TERSE
diff --git a/py/runtime0.h b/py/runtime0.h
@@ -76,6 +76,8 @@ typedef enum {
     MP_UNARY_OP_HASH, // __hash__; must return a small int
     MP_UNARY_OP_ABS, // __abs__
     MP_UNARY_OP_INT, // __int__
+    MP_UNARY_OP_FLOAT_MAYBE, // __float__
+    MP_UNARY_OP_COMPLEX_MAYBE, // __complex__
     MP_UNARY_OP_SIZEOF, // for sys.getsizeof()
 } mp_unary_op_t;
 
diff --git a/tests/cpydiff/types_float_implicit_conversion.py b/tests/cpydiff/types_float_implicit_conversion.py
@@ -0,0 +1,14 @@
+"""
+categories: Types,float
+description: uPy allows implicit conversion of objects in maths operations while CPython does not.
+cause: Unknown
+workaround: Objects should be wrapped in `float(obj)` for compatibility with CPython.
+"""
+
+
+class Test:
+    def __float__(self):
+        return 0.5
+
+
+print(2.0 * Test())
diff --git a/tests/float/complex_dunder.py b/tests/float/complex_dunder.py
@@ -0,0 +1,40 @@
+# test __complex__ function support
+
+
+class TestComplex:
+    def __complex__(self):
+        return 1j + 10
+
+
+class TestStrComplex:
+    def __complex__(self):
+        return "a"
+
+
+class TestNonComplex:
+    def __complex__(self):
+        return 6
+
+
+class Test:
+    pass
+
+
+print(complex(TestComplex()))
+
+try:
+    print(complex(TestStrComplex()))
+except TypeError:
+    print("TypeError")
+
+
+try:
+    print(complex(TestNonComplex()))
+except TypeError:
+    print("TypeError")
+
+
+try:
+    print(complex(Test()))
+except TypeError:
+    print("TypeError")
diff --git a/tests/float/float_dunder.py b/tests/float/float_dunder.py
@@ -0,0 +1,42 @@
+# test __float__ function support
+
+
+class TestFloat:
+    def __float__(self):
+        return 10.0
+
+
+class TestStrFloat:
+    def __float__(self):
+        return "a"
+
+
+class TestNonFloat:
+    def __float__(self):
+        return 6
+
+
+class Test:
+    pass
+
+
+print("%.1f" % float(TestFloat()))
+print("%.1f" % TestFloat())
+
+
+try:
+    print(float(TestStrFloat()))
+except TypeError:
+    print("TypeError")
+
+
+try:
+    print(float(TestNonFloat()))
+except TypeError:
+    print("TypeError")
+
+
+try:
+    print(float(Test()))
+except TypeError:
+    print("TypeError")
diff --git a/tests/run-tests.py b/tests/run-tests.py
@@ -525,6 +525,7 @@ def run_tests(pyb, tests, args, result_dir, num_threads=1):
         skip_tests.add("float/int_big_float.py")
         skip_tests.add("float/true_value.py")
         skip_tests.add("float/types.py")
+        skip_tests.add("float/complex_dunder.py")
 
     if not has_coverage:
         skip_tests.add("cmdline/cmd_parsetree.py")
