diff options
| -rw-r--r-- | py/objstr.c | 86 | ||||
| -rw-r--r-- | py/runtime.c | 15 | ||||
| -rw-r--r-- | tests/basics/string_replace.py | 8 | ||||
| -rw-r--r-- | tests/basics/string_strip.py | 6 |
4 files changed, 107 insertions, 8 deletions
diff --git a/py/objstr.c b/py/objstr.c index c30da00d71..fa658ac4d7 100644 --- a/py/objstr.c +++ b/py/objstr.c @@ -408,12 +408,97 @@ mp_obj_t str_format(uint n_args, const mp_obj_t *args) { return s; } +mp_obj_t str_replace(uint n_args, const mp_obj_t *args) { + assert(MP_OBJ_IS_STR(args[0])); + assert(MP_OBJ_IS_STR(args[1])); + assert(MP_OBJ_IS_STR(args[2])); + + machine_int_t max_rep = 0; + if (n_args == 4) { + assert(MP_OBJ_IS_SMALL_INT(args[3])); + max_rep = MP_OBJ_SMALL_INT_VALUE(args[3]); + if (max_rep == 0) { + return args[0]; + } else if (max_rep < 0) { + max_rep = 0; + } + } + + // if max_rep is still 0 by this point we will need to do all possible replacements + + GET_STR_DATA_LEN(args[0], str, str_len); + GET_STR_DATA_LEN(args[1], old, old_len); + GET_STR_DATA_LEN(args[2], new, new_len); + + // old won't exist in str if it's longer, so nothing to replace + if (old_len > str_len) { + return args[0]; + } + + // data for the replaced string + byte *data = NULL; + mp_obj_t replaced_str = MP_OBJ_NULL; + + // do 2 passes over the string: + // first pass computes the required length of the replaced string + // second pass does the replacements + for (;;) { + machine_uint_t replaced_str_index = 0; + machine_uint_t num_replacements_done = 0; + const byte *old_occurrence; + const byte *offset_ptr = str; + machine_uint_t offset_num = 0; + while ((old_occurrence = find_subbytes(offset_ptr, str_len - offset_num, old, old_len)) != NULL) { + // copy from just after end of last occurrence of to-be-replaced string to right before start of next occurrence + if (data != NULL) { + memcpy(data + replaced_str_index, offset_ptr, old_occurrence - offset_ptr); + } + replaced_str_index += old_occurrence - offset_ptr; + // copy the replacement string + if (data != NULL) { + memcpy(data + replaced_str_index, new, new_len); + } + replaced_str_index += new_len; + offset_ptr = old_occurrence + old_len; + offset_num = offset_ptr - str; + + num_replacements_done++; + if (max_rep != 0 && num_replacements_done == max_rep){ + break; + } + } + + // copy from just after end of last occurrence of to-be-replaced string to end of old string + if (data != NULL) { + memcpy(data + replaced_str_index, offset_ptr, str_len - offset_num); + } + replaced_str_index += str_len - offset_num; + + if (data == NULL) { + // first pass + if (num_replacements_done == 0) { + // no substr found, return original string + return args[0]; + } else { + // substr found, allocate new string + replaced_str = mp_obj_str_builder_start(mp_obj_get_type(args[0]), replaced_str_index, &data); + } + } else { + // second pass, we are done + break; + } + } + + return mp_obj_str_builder_end(replaced_str); +} + static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_find_obj, 2, 4, str_find); 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(str_format_obj, 1, str_format); +static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_replace_obj, 3, 4, str_replace); static const mp_method_t str_type_methods[] = { { "find", &str_find_obj }, @@ -422,6 +507,7 @@ static const mp_method_t str_type_methods[] = { { "startswith", &str_startswith_obj }, { "strip", &str_strip_obj }, { "format", &str_format_obj }, + { "replace", &str_replace_obj }, { NULL, NULL }, // end-of-list sentinel }; diff --git a/py/runtime.c b/py/runtime.c index ea96ca9f7a..ee8d720c22 100644 --- a/py/runtime.c +++ b/py/runtime.c @@ -737,15 +737,14 @@ mp_obj_t rt_call_function_n_kw(mp_obj_t fun_in, uint n_args, uint n_kw, const mp DEBUG_OP_printf("calling function %p(n_args=%d, n_kw=%d, args=%p)\n", fun_in, n_args, n_kw, args); - if (MP_OBJ_IS_SMALL_INT(fun_in)) { - nlr_jump(mp_obj_new_exception_msg(MP_QSTR_TypeError, "'int' object is not callable")); + // get the type + mp_obj_type_t *type = mp_obj_get_type(fun_in); + + // do the call + if (type->call != NULL) { + return type->call(fun_in, n_args, n_kw, args); } else { - mp_obj_base_t *fun = fun_in; - if (fun->type->call != NULL) { - return fun->type->call(fun_in, n_args, n_kw, args); - } else { - nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_TypeError, "'%s' object is not callable", fun->type->name)); - } + nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_TypeError, "'%s' object is not callable", type->name)); } } diff --git a/tests/basics/string_replace.py b/tests/basics/string_replace.py new file mode 100644 index 0000000000..a1d0f973bb --- /dev/null +++ b/tests/basics/string_replace.py @@ -0,0 +1,8 @@ +print("".replace("a", "b")) +print("aaa".replace("a", "b", 0)) +print("aaa".replace("a", "b", -5)) +print("asdfasdf".replace("a", "b")) +print("aabbaabbaabbaa".replace("aa", "cc", 3)) +print("a".replace("aa", "bb")) +print("testingtesting".replace("ing", "")) +print("testINGtesting".replace("ing", "ING!")) diff --git a/tests/basics/string_strip.py b/tests/basics/string_strip.py new file mode 100644 index 0000000000..518dfd66ec --- /dev/null +++ b/tests/basics/string_strip.py @@ -0,0 +1,6 @@ +print("".strip()) +print(" \t\n\r\v\f".strip()) +print(" T E S T".strip()) +print("abcabc".strip("ce")) +print("aaa".strip("b")) +print("abc efg ".strip("g a")) |