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

5 files changed, 33 insertions, 3 deletions

diff --git a/py/obj.c b/py/obj.c
index 51f6d85de..5a05ea58c 100644
--- 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
index 157617e15..3c4cb6614 100644
--- 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
index 37c1e3bd2..fe1918bd3 100644
--- 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
index e6d8c6807..2c3b3ddde 100644
--- 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
index e6eeff97d..c82a4717f 100644
--- 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;