Merge remote-tracking branch 'upstream/master' into list_insert

1 files changed, 25 insertions, 3 deletions

diff --git a/py/objstr.c b/py/objstr.c
index 48abf4951d..59547e3cd6 100644
--- a/py/objstr.c
+++ b/py/objstr.c
@@ -27,9 +27,31 @@ mp_obj_t str_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
     const char *lhs_str = qstr_str(lhs->qstr);
     switch (op) {
         case RT_BINARY_OP_SUBSCR:
-            // string access
-            // XXX a massive hack!
-            return mp_obj_new_int(lhs_str[mp_obj_get_int(rhs_in)]);
+            // 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)) {
+                // TODO: This implements byte string access for single index so far
+                // TODO: Handle negative indexes.
+                return mp_obj_new_int(lhs_str[mp_obj_get_int(rhs_in)]);
+#if MICROPY_ENABLE_SLICE
+            } else if (MP_OBJ_IS_TYPE(rhs_in, &slice_type)) {
+                machine_int_t start, stop, step;
+                mp_obj_slice_get(rhs_in, &start, &stop, &step);
+                assert(step == 1);
+                int len = strlen(lhs_str);
+                if (start < 0) {
+                    start = len + start;
+                }
+                if (stop <= 0) {
+                    stop = len + stop;
+                }
+                return mp_obj_new_str(qstr_from_strn_copy(lhs_str + start, stop - start));
+#endif
+            } else {
+                // Message doesn't match CPython, but we don't have so much bytes as they
+                // to spend them on verbose wording
+                nlr_jump(mp_obj_new_exception_msg(rt_q_TypeError, "index must be int"));
+            }
 
         case RT_BINARY_OP_ADD:
         case RT_BINARY_OP_INPLACE_ADD: