diff options
Diffstat (limited to 'py/objstr.c')
| -rw-r--r-- | py/objstr.c | 71 |
1 files changed, 60 insertions, 11 deletions
diff --git a/py/objstr.c b/py/objstr.c index ea4f5ead24..0621a8df75 100644 --- a/py/objstr.c +++ b/py/objstr.c @@ -22,10 +22,14 @@ static mp_obj_t mp_obj_new_str_iterator(mp_obj_str_t *str, int cur); /******************************************************************************/ /* str */ -void str_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in) { +void str_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in, mp_print_kind_t kind) { mp_obj_str_t *self = self_in; - // TODO need to escape chars etc - print(env, "'%s'", qstr_str(self->qstr)); + if (kind == PRINT_STR) { + print(env, "%s", qstr_str(self->qstr)); + } else { + // TODO need to escape chars etc + print(env, "'%s'", qstr_str(self->qstr)); + } } mp_obj_t str_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) { @@ -85,6 +89,15 @@ mp_obj_t str_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) { return mp_obj_new_str(qstr_from_str_take(val, alloc_len)); } break; + case RT_COMPARE_OP_IN: + case RT_COMPARE_OP_NOT_IN: + /* NOTE `a in b` is `b.__contains__(a)` */ + if (MP_OBJ_IS_TYPE(rhs_in, &str_type)) { + const char *rhs_str = qstr_str(((mp_obj_str_t*)rhs_in)->qstr); + /* FIXME \0 in strs */ + return MP_BOOL((op == RT_COMPARE_OP_IN) ^ (strstr(lhs_str, rhs_str) == NULL)); + } + break; } return MP_OBJ_NULL; // op not supported @@ -156,6 +169,46 @@ static bool chr_in_str(const char* const str, const size_t str_len, const char c return false; } +static mp_obj_t str_find(int n_args, const mp_obj_t *args) { + assert(2 <= n_args && n_args <= 4); + assert(MP_OBJ_IS_TYPE(args[0], &str_type)); + if (!MP_OBJ_IS_TYPE(args[1], &str_type)) { + nlr_jump(mp_obj_new_exception_msg_1_arg( + MP_QSTR_TypeError, + "Can't convert '%s' object to str implicitly", + mp_obj_get_type_str(args[1]))); + } + + const char* haystack = qstr_str(((mp_obj_str_t*)args[0])->qstr); + const char* needle = qstr_str(((mp_obj_str_t*)args[1])->qstr); + + size_t haystack_len = strlen(haystack); + size_t needle_len = strlen(needle); + + size_t start = 0; + size_t end = haystack_len; + /* TODO use a non-exception-throwing mp_get_index */ + if (n_args >= 3 && args[2] != mp_const_none) { + start = mp_get_index(&str_type, haystack_len, args[2]); + } + if (n_args >= 4 && args[3] != mp_const_none) { + end = mp_get_index(&str_type, haystack_len, args[3]); + } + + char *p = strstr(haystack + start, needle); + if (p == NULL) { + // not found + return MP_OBJ_NEW_SMALL_INT(-1); + } else { + // found + machine_int_t pos = p - haystack; + if (pos + needle_len > end) { + pos = -1; + } + return MP_OBJ_NEW_SMALL_INT(pos); + } +} + mp_obj_t str_strip(int n_args, const mp_obj_t *args) { assert(1 <= n_args && n_args <= 2); assert(MP_OBJ_IS_TYPE(args[0], &str_type)); @@ -205,13 +258,6 @@ mp_obj_t str_strip(int n_args, const mp_obj_t *args) { return mp_obj_new_str(qstr_from_str_take(stripped_str, stripped_len + 1)); } -void vstr_printf_wrapper(void *env, const char *fmt, ...) { - va_list args; - va_start(args, fmt); - vstr_vprintf(env, fmt, args); - va_end(args); -} - mp_obj_t str_format(int n_args, const mp_obj_t *args) { assert(MP_OBJ_IS_TYPE(args[0], &str_type)); mp_obj_str_t *self = args[0]; @@ -228,7 +274,8 @@ mp_obj_t str_format(int n_args, const mp_obj_t *args) { if (arg_i >= n_args) { nlr_jump(mp_obj_new_exception_msg(MP_QSTR_IndexError, "tuple index out of range")); } - mp_obj_print_helper(vstr_printf_wrapper, vstr, args[arg_i]); + // TODO: may be PRINT_REPR depending on formatting code + mp_obj_print_helper((void (*)(void*, const char*, ...))vstr_printf, vstr, args[arg_i], PRINT_STR); arg_i++; } } else { @@ -239,11 +286,13 @@ mp_obj_t str_format(int n_args, const mp_obj_t *args) { return mp_obj_new_str(qstr_from_str_take(vstr->buf, vstr->alloc)); } +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_strip_obj, 1, 2, str_strip); static MP_DEFINE_CONST_FUN_OBJ_VAR(str_format_obj, 1, str_format); static const mp_method_t str_type_methods[] = { + { "find", &str_find_obj }, { "join", &str_join_obj }, { "strip", &str_strip_obj }, { "format", &str_format_obj }, |