diff options
Diffstat (limited to 'py/objstr.c')
| -rw-r--r-- | py/objstr.c | 43 |
1 files changed, 28 insertions, 15 deletions
diff --git a/py/objstr.c b/py/objstr.c index 5e9f0c76b6..519ff464d6 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]); @@ -329,17 +327,19 @@ STATIC mp_obj_t str_join(mp_obj_t self_in, mp_obj_t arg) { mp_obj_t *seq_items; if (MP_OBJ_IS_TYPE(arg, &mp_type_tuple)) { mp_obj_tuple_get(arg, &seq_len, &seq_items); - } else if (MP_OBJ_IS_TYPE(arg, &mp_type_list)) { - mp_obj_list_get(arg, &seq_len, &seq_items); } else { - goto bad_arg; + if (!MP_OBJ_IS_TYPE(arg, &mp_type_list)) { + // arg is not a list, try to convert it to one + arg = mp_type_list.make_new((mp_obj_t)&mp_type_list, 1, 0, &arg); + } + mp_obj_list_get(arg, &seq_len, &seq_items); } // count required length int required_len = 0; for (int i = 0; i < seq_len; i++) { if (!MP_OBJ_IS_STR(seq_items[i])) { - goto bad_arg; + nlr_raise(mp_obj_new_exception_msg(&mp_type_TypeError, "join expected a list of str's")); } if (i > 0) { required_len += sep_len; @@ -363,9 +363,6 @@ STATIC mp_obj_t str_join(mp_obj_t self_in, mp_obj_t arg) { // return joined string return mp_obj_str_builder_end(joined_str); - -bad_arg: - nlr_raise(mp_obj_new_exception_msg(&mp_type_TypeError, "?str.join expecting a list of str's")); } #define is_ws(c) ((c) == ' ' || (c) == '\t') @@ -441,7 +438,7 @@ STATIC mp_obj_t str_split(uint n_args, const mp_obj_t *args) { return res; } -STATIC mp_obj_t str_finder(uint n_args, const mp_obj_t *args, machine_int_t direction) { +STATIC mp_obj_t str_finder(uint n_args, const mp_obj_t *args, machine_int_t direction, bool is_index) { assert(2 <= n_args && n_args <= 4); assert(MP_OBJ_IS_STR(args[0])); assert(MP_OBJ_IS_STR(args[1])); @@ -461,7 +458,11 @@ STATIC mp_obj_t str_finder(uint n_args, const mp_obj_t *args, machine_int_t dire const byte *p = find_subbytes(haystack + start, end - start, needle, needle_len, direction); if (p == NULL) { // not found - return MP_OBJ_NEW_SMALL_INT(-1); + if (is_index) { + nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "substring not found")); + } else { + return MP_OBJ_NEW_SMALL_INT(-1); + } } else { // found return MP_OBJ_NEW_SMALL_INT(p - haystack); @@ -469,11 +470,19 @@ STATIC mp_obj_t str_finder(uint n_args, const mp_obj_t *args, machine_int_t dire } STATIC mp_obj_t str_find(uint n_args, const mp_obj_t *args) { - return str_finder(n_args, args, 1); + return str_finder(n_args, args, 1, false); } STATIC mp_obj_t str_rfind(uint n_args, const mp_obj_t *args) { - return str_finder(n_args, args, -1); + return str_finder(n_args, args, -1, false); +} + +STATIC mp_obj_t str_index(uint n_args, const mp_obj_t *args) { + return str_finder(n_args, args, 1, true); +} + +STATIC mp_obj_t str_rindex(uint n_args, const mp_obj_t *args) { + return str_finder(n_args, args, -1, true); } // TODO: (Much) more variety in args @@ -1128,7 +1137,7 @@ STATIC mp_obj_t str_replace(uint n_args, const mp_obj_t *args) { } } - // if max_rep is still 0 by this point we will need to do all possible replacements + // if max_rep is still -1 by this point we will need to do all possible replacements // check argument types @@ -1307,6 +1316,8 @@ STATIC machine_int_t str_get_buffer(mp_obj_t self_in, buffer_info_t *bufinfo, in STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_find_obj, 2, 4, str_find); STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_rfind_obj, 2, 4, str_rfind); +STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_index_obj, 2, 4, str_index); +STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_rindex_obj, 2, 4, str_rindex); STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_join_obj, str_join); STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_split_obj, 1, 3, str_split); STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_startswith_obj, str_startswith); @@ -1320,6 +1331,8 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_rpartition_obj, str_rpartition); STATIC const mp_map_elem_t str_locals_dict_table[] = { { MP_OBJ_NEW_QSTR(MP_QSTR_find), (mp_obj_t)&str_find_obj }, { MP_OBJ_NEW_QSTR(MP_QSTR_rfind), (mp_obj_t)&str_rfind_obj }, + { MP_OBJ_NEW_QSTR(MP_QSTR_index), (mp_obj_t)&str_index_obj }, + { MP_OBJ_NEW_QSTR(MP_QSTR_rindex), (mp_obj_t)&str_rindex_obj }, { MP_OBJ_NEW_QSTR(MP_QSTR_join), (mp_obj_t)&str_join_obj }, { MP_OBJ_NEW_QSTR(MP_QSTR_split), (mp_obj_t)&str_split_obj }, { MP_OBJ_NEW_QSTR(MP_QSTR_startswith), (mp_obj_t)&str_startswith_obj }, |