diff options
Diffstat (limited to 'py/objstr.c')
| -rw-r--r-- | py/objstr.c | 39 |
1 files changed, 36 insertions, 3 deletions
diff --git a/py/objstr.c b/py/objstr.c index b0d5cba6f1..6819a4ad68 100644 --- a/py/objstr.c +++ b/py/objstr.c @@ -493,7 +493,9 @@ STATIC mp_obj_t str_startswith(mp_obj_t self_in, mp_obj_t arg) { return MP_BOOL(memcmp(str, prefix, prefix_len) == 0); } -STATIC mp_obj_t str_strip(uint n_args, const mp_obj_t *args) { +enum { LSTRIP, RSTRIP, STRIP }; + +STATIC mp_obj_t str_uni_strip(int type, uint n_args, const mp_obj_t *args) { assert(1 <= n_args && n_args <= 2); assert(MP_OBJ_IS_STR(args[0])); @@ -516,14 +518,29 @@ STATIC mp_obj_t str_strip(uint n_args, const mp_obj_t *args) { machine_uint_t first_good_char_pos = 0; bool first_good_char_pos_set = false; machine_uint_t last_good_char_pos = 0; - for (machine_uint_t i = 0; i < orig_str_len; i++) { + machine_uint_t i = 0; + machine_int_t delta = 1; + if (type == RSTRIP) { + i = orig_str_len - 1; + delta = -1; + } + for (machine_uint_t len = orig_str_len; len > 0; len--) { if (find_subbytes(chars_to_del, chars_to_del_len, &orig_str[i], 1, 1) == NULL) { - last_good_char_pos = i; if (!first_good_char_pos_set) { first_good_char_pos = i; + if (type == LSTRIP) { + last_good_char_pos = orig_str_len - 1; + break; + } else if (type == RSTRIP) { + first_good_char_pos = 0; + last_good_char_pos = i; + break; + } first_good_char_pos_set = true; } + last_good_char_pos = i; } + i += delta; } if (first_good_char_pos == 0 && last_good_char_pos == 0) { @@ -537,6 +554,18 @@ STATIC mp_obj_t str_strip(uint n_args, const mp_obj_t *args) { return mp_obj_new_str(orig_str + first_good_char_pos, stripped_len, false); } +STATIC mp_obj_t str_strip(uint n_args, const mp_obj_t *args) { + return str_uni_strip(STRIP, n_args, args); +} + +STATIC mp_obj_t str_lstrip(uint n_args, const mp_obj_t *args) { + return str_uni_strip(LSTRIP, n_args, args); +} + +STATIC mp_obj_t str_rstrip(uint n_args, const mp_obj_t *args) { + return str_uni_strip(RSTRIP, n_args, args); +} + // Takes an int arg, but only parses unsigned numbers, and only changes // *num if at least one digit was parsed. static int str_to_int(const char *str, int *num) { @@ -1354,6 +1383,8 @@ 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); STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_strip_obj, 1, 2, str_strip); +STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_lstrip_obj, 1, 2, str_lstrip); +STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_rstrip_obj, 1, 2, str_rstrip); STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(str_format_obj, 1, mp_obj_str_format); STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_replace_obj, 3, 4, str_replace); STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_count_obj, 2, 4, str_count); @@ -1373,6 +1404,8 @@ STATIC const mp_map_elem_t str_locals_dict_table[] = { { 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 }, { MP_OBJ_NEW_QSTR(MP_QSTR_strip), (mp_obj_t)&str_strip_obj }, + { MP_OBJ_NEW_QSTR(MP_QSTR_lstrip), (mp_obj_t)&str_lstrip_obj }, + { MP_OBJ_NEW_QSTR(MP_QSTR_rstrip), (mp_obj_t)&str_rstrip_obj }, { MP_OBJ_NEW_QSTR(MP_QSTR_format), (mp_obj_t)&str_format_obj }, { MP_OBJ_NEW_QSTR(MP_QSTR_replace), (mp_obj_t)&str_replace_obj }, { MP_OBJ_NEW_QSTR(MP_QSTR_count), (mp_obj_t)&str_count_obj }, |