py: Add mp_obj_is_integer; make mp_get_index check for long int.

mp_obj_is_integer should be used to check if an object is of integral
type.  It returns true for bool, small int and long int.

3 files changed, 9 insertions, 5 deletions

diff --git a/py/obj.c b/py/obj.c
index d1db53690f..844ec41216 100644
--- a/py/obj.c
+++ b/py/obj.c
@@ -99,6 +99,11 @@ int mp_obj_is_true(mp_obj_t arg) {
     }
 }
 
+// returns true if o_in is bool, small int, or long int
+bool mp_obj_is_integer(mp_obj_t o_in) {
+    return MP_OBJ_IS_INT(o_in) || MP_OBJ_IS_TYPE(o_in, &mp_type_bool);
+}
+
 bool mp_obj_is_callable(mp_obj_t o_in) {
     return mp_obj_get_type(o_in)->call != NULL;
 }
@@ -285,8 +290,8 @@ void mp_obj_get_array_fixed_n(mp_obj_t o, uint len, mp_obj_t **items) {
 // is_slice determines whether the index is a slice index
 uint mp_get_index(const mp_obj_type_t *type, machine_uint_t len, mp_obj_t index, bool is_slice) {
     int i;
-    if (MP_OBJ_IS_SMALL_INT(index)) {
-        i = MP_OBJ_SMALL_INT_VALUE(index);
+    if (MP_OBJ_IS_INT(index)) {
+        i = mp_obj_int_get_checked(index);
     } else if (MP_OBJ_IS_TYPE(index, &mp_type_bool)) {
         i = (index == mp_const_true ? 1 : 0);
     } else {
diff --git a/py/obj.h b/py/obj.h
index ab1685d4d0..b944d2fcb5 100644
--- a/py/obj.h
+++ b/py/obj.h
@@ -369,6 +369,7 @@ void mp_obj_print(mp_obj_t o, mp_print_kind_t kind);
 void mp_obj_print_exception(mp_obj_t exc);
 
 int mp_obj_is_true(mp_obj_t arg);
+bool mp_obj_is_integer(mp_obj_t o_in); // returns true if o_in is bool, small int, or long int
 bool mp_obj_is_callable(mp_obj_t o_in);
 machine_int_t mp_obj_hash(mp_obj_t o_in);
 bool mp_obj_equal(mp_obj_t o1, mp_obj_t o2);
diff --git a/py/objstr.c b/py/objstr.c
index 06bc98a723..14f32459d1 100644
--- a/py/objstr.c
+++ b/py/objstr.c
@@ -221,9 +221,7 @@ STATIC mp_obj_t str_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
     GET_STR_DATA_LEN(lhs_in, lhs_data, lhs_len);
     switch (op) {
         case MP_BINARY_OP_SUBSCR:
-            // TODO: need predicate to check for int-like type (bools are such for example)
-            // ["no", "yes"][1 == 2] is common idiom
-            if (MP_OBJ_IS_SMALL_INT(rhs_in)) {
+            if (mp_obj_is_integer(rhs_in)) {
                 uint index = mp_get_index(mp_obj_get_type(lhs_in), lhs_len, rhs_in, false);
                 if (MP_OBJ_IS_TYPE(lhs_in, &mp_type_bytes)) {
                     return MP_OBJ_NEW_SMALL_INT((mp_small_int_t)lhs_data[index]);