Skip to content

Commit a9ddd6d

Browse files
dpgeorgedpgeorge
committed
py: Simplify and improve mp_get_index.
It has (again) a fast path for ints, and a simplified "slow" path for everything else. Also simplify the way str indexing is done (now matches tuple and list).
1 parent 5f3fe3a commit a9ddd6d

2 files changed

Lines changed: 13 additions & 18 deletions

File tree

‎py/obj.c‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -277,12 +277,10 @@ void mp_obj_get_array_fixed_n(mp_obj_t o, uint len, mp_obj_t **items) {
277277

278278
// is_slice determines whether the index is a slice index
279279
uint mp_get_index(const mp_obj_type_t *type, machine_uint_t len, mp_obj_t index, bool is_slice) {
280-
int i;
281-
if (MP_OBJ_IS_INT(index)) {
282-
i = mp_obj_int_get_checked(index);
283-
} else if (MP_OBJ_IS_TYPE(index, &mp_type_bool)) {
284-
i = (index == mp_const_true ? 1 : 0);
285-
} else {
280+
machine_int_t i;
281+
if (MP_OBJ_IS_SMALL_INT(index)) {
282+
i = MP_OBJ_SMALL_INT_VALUE(index);
283+
} else if (!mp_obj_get_int_maybe(index, &i)) {
286284
nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_TypeError, "%s indices must be integers, not %s", qstr_str(type->name), mp_obj_get_type_str(index)));
287285
}
288286

‎py/objstr.c‎

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -220,27 +220,24 @@ STATIC const byte *find_subbytes(const byte *haystack, machine_uint_t hlen, cons
220220
STATIC mp_obj_t str_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
221221
GET_STR_DATA_LEN(lhs_in, lhs_data, lhs_len);
222222
switch (op) {
223-
case MP_BINARY_OP_SUBSCR:
224-
if (mp_obj_is_integer(rhs_in)) {
225-
uint index = mp_get_index(mp_obj_get_type(lhs_in), lhs_len, rhs_in, false);
226-
if (MP_OBJ_IS_TYPE(lhs_in, &mp_type_bytes)) {
227-
return MP_OBJ_NEW_SMALL_INT((mp_small_int_t)lhs_data[index]);
228-
} else {
229-
return mp_obj_new_str(lhs_data + index, 1, true);
230-
}
223+
case MP_BINARY_OP_SUBSCR: {
231224
#if MICROPY_ENABLE_SLICE
232-
} else if (MP_OBJ_IS_TYPE(rhs_in, &mp_type_slice)) {
225+
if (MP_OBJ_IS_TYPE(rhs_in, &mp_type_slice)) {
233226
machine_uint_t start, stop;
234227
if (!m_seq_get_fast_slice_indexes(lhs_len, rhs_in, &start, &stop)) {
235228
assert(0);
236229
}
237230
return mp_obj_new_str(lhs_data + start, stop - start, false);
231+
}
238232
#endif
233+
mp_obj_type_t *type = mp_obj_get_type(lhs_in);
234+
uint index = mp_get_index(type, lhs_len, rhs_in, false);
235+
if (type == &mp_type_bytes) {
236+
return MP_OBJ_NEW_SMALL_INT((mp_small_int_t)lhs_data[index]);
239237
} else {
240-
// Message doesn't match CPython, but we don't have so much bytes as they
241-
// to spend them on verbose wording
242-
nlr_raise(mp_obj_new_exception_msg(&mp_type_TypeError, "index must be int"));
238+
return mp_obj_new_str(lhs_data + index, 1, true);
243239
}
240+
}
244241

245242
case MP_BINARY_OP_ADD:
246243
case MP_BINARY_OP_INPLACE_ADD:

0 commit comments

Comments
 (0)