From b1d08726eeeadf0b91358cc8c46e156695b6a9c0 Mon Sep 17 00:00:00 2001
From: Paul Sokolovsky <pfalcon@users.sourceforge.net>
Date: Wed, 29 Aug 2018 17:59:36 +0300
Subject: [PATCH] py/obj: Add support for __int__ special method.

Based on the discussion, this special method is available unconditionally,
as converting to int is a common operation.
---
 py/obj.c      |  8 ++------
 py/objint.c   |  3 +--
 py/objtype.c  | 19 ++++++++++++++++---
 py/runtime.c  | 20 ++++++++++++++++----
 py/runtime0.h |  1 +
 5 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/py/obj.c b/py/obj.c
index 5eb2b094e..47b7d15ae 100644
--- a/py/obj.c
+++ b/py/obj.c
@@ -235,12 +235,8 @@ mp_int_t mp_obj_get_int(mp_const_obj_t arg) {
     } else if (MP_OBJ_IS_TYPE(arg, &mp_type_int)) {
         return mp_obj_int_get_checked(arg);
     } else {
-        if (MICROPY_ERROR_REPORTING == MICROPY_ERROR_REPORTING_TERSE) {
-            mp_raise_TypeError("can't convert to int");
-        } else {
-            nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_TypeError,
-                "can't convert %s to int", mp_obj_get_type_str(arg)));
-        }
+        mp_obj_t res = mp_unary_op(MP_UNARY_OP_INT, (mp_obj_t)arg);
+        return mp_obj_int_get_checked(res);
     }
 }
 
diff --git a/py/objint.c b/py/objint.c
index cd8f20c34..d5d74dd55 100644
--- a/py/objint.c
+++ b/py/objint.c
@@ -62,8 +62,7 @@ STATIC mp_obj_t mp_obj_int_make_new(const mp_obj_type_t *type_in, size_t n_args,
                 return mp_obj_new_int_from_float(mp_obj_float_get(args[0]));
 #endif
             } else {
-                // try to convert to small int (eg from bool)
-                return MP_OBJ_NEW_SMALL_INT(mp_obj_get_int(args[0]));
+                return mp_unary_op(MP_UNARY_OP_INT, args[0]);
             }
 
         case 2:
diff --git a/py/objtype.c b/py/objtype.c
index 0881ae33f..fec73f16e 100644
--- a/py/objtype.c
+++ b/py/objtype.c
@@ -376,6 +376,7 @@ const byte mp_unary_op_method_name[MP_UNARY_OP_NUM_RUNTIME] = {
     [MP_UNARY_OP_BOOL] = MP_QSTR___bool__,
     [MP_UNARY_OP_LEN] = MP_QSTR___len__,
     [MP_UNARY_OP_HASH] = MP_QSTR___hash__,
+    [MP_UNARY_OP_INT] = MP_QSTR___int__,
     #if MICROPY_PY_ALL_SPECIAL_METHODS
     [MP_UNARY_OP_POSITIVE] = MP_QSTR___pos__,
     [MP_UNARY_OP_NEGATIVE] = MP_QSTR___neg__,
@@ -421,9 +422,21 @@ STATIC mp_obj_t instance_unary_op(mp_unary_op_t op, mp_obj_t self_in) {
         return mp_unary_op(op, self->subobj[0]);
     } else if (member[0] != MP_OBJ_NULL) {
         mp_obj_t val = mp_call_function_1(member[0], self_in);
-        // __hash__ must return a small int
-        if (op == MP_UNARY_OP_HASH) {
-            val = MP_OBJ_NEW_SMALL_INT(mp_obj_get_int_truncated(val));
+
+        switch (op) {
+            case MP_UNARY_OP_HASH:
+                // __hash__ must return a small int
+                val = MP_OBJ_NEW_SMALL_INT(mp_obj_get_int_truncated(val));
+                break;
+            case MP_UNARY_OP_INT:
+                // Must return int
+                if (!MP_OBJ_IS_INT(val)) {
+                    mp_raise_TypeError(NULL);
+                }
+                break;
+            default:
+                // No need to do anything
+                ;
         }
         return val;
     } else {
diff --git a/py/runtime.c b/py/runtime.c
index 33c4c1822..e512b8b0d 100644
--- a/py/runtime.c
+++ b/py/runtime.c
@@ -228,6 +228,7 @@ mp_obj_t mp_unary_op(mp_unary_op_t op, mp_obj_t arg) {
             case MP_UNARY_OP_HASH:
                 return arg;
             case MP_UNARY_OP_POSITIVE:
+            case MP_UNARY_OP_INT:
                 return arg;
             case MP_UNARY_OP_NEGATIVE:
                 // check for overflow
@@ -265,12 +266,23 @@ mp_obj_t mp_unary_op(mp_unary_op_t op, mp_obj_t arg) {
                 return result;
             }
         }
+        // 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) {
-            mp_raise_TypeError("unsupported type for operator");
+            if (op == MP_UNARY_OP_INT) {
+                mp_raise_TypeError("can't convert to int");
+            } else {
+                mp_raise_TypeError("unsupported type for operator");
+            }
         } else {
-            nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_TypeError,
-                "unsupported type for %q: '%s'",
-                mp_unary_op_method_name[op], mp_obj_get_type_str(arg)));
+            if (op == MP_UNARY_OP_INT) {
+                nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_TypeError,
+                    "can't convert %s to int", mp_obj_get_type_str(arg)));
+            } else {
+                nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_TypeError,
+                    "unsupported type for %q: '%s'",
+                    mp_unary_op_method_name[op], mp_obj_get_type_str(arg)));
+            }
         }
     }
 }
diff --git a/py/runtime0.h b/py/runtime0.h
index 56cc6cfd3..efd439196 100644
--- a/py/runtime0.h
+++ b/py/runtime0.h
@@ -61,6 +61,7 @@ typedef enum {
     MP_UNARY_OP_LEN, // __len__
     MP_UNARY_OP_HASH, // __hash__; must return a small int
     MP_UNARY_OP_ABS, // __abs__
+    MP_UNARY_OP_INT, // __int__
     MP_UNARY_OP_SIZEOF, // for sys.getsizeof()
 
     MP_UNARY_OP_NUM_RUNTIME,
-- 
GitLab