Skip to content

Commit 4c31655

Browse files
committed
Implement str.split(None).
Note that splitting by explicit string is not implemented so far.
1 parent 7380a83 commit 4c31655

2 files changed

Lines changed: 47 additions & 0 deletions

File tree

py/objstr.c

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,44 @@ mp_obj_t str_join(mp_obj_t self_in, mp_obj_t arg) {
175175
nlr_jump(mp_obj_new_exception_msg(MP_QSTR_TypeError, "?str.join expecting a list of str's"));
176176
}
177177

178+
#define is_ws(c) ((c) == ' ' || (c) == '\t')
179+
180+
static mp_obj_t str_split(uint n_args, const mp_obj_t *args) {
181+
int splits = -1;
182+
mp_obj_t sep = mp_const_none;
183+
if (n_args > 1) {
184+
sep = args[1];
185+
if (n_args > 2) {
186+
splits = MP_OBJ_SMALL_INT_VALUE(args[2]);
187+
}
188+
}
189+
assert(sep == mp_const_none);
190+
mp_obj_t res = mp_obj_new_list(0, NULL);
191+
const char *s = qstr_str(mp_obj_str_get(args[0]));
192+
const char *start;
193+
194+
// Initial whitespace is not counted as split, so we pre-do it
195+
while (is_ws(*s)) s++;
196+
while (*s && splits != 0) {
197+
start = s;
198+
while (*s != 0 && !is_ws(*s)) s++;
199+
rt_list_append(res, MP_OBJ_NEW_QSTR(qstr_from_strn_copy(start, s - start)));
200+
if (*s == 0) {
201+
break;
202+
}
203+
while (is_ws(*s)) s++;
204+
if (splits > 0) {
205+
splits--;
206+
}
207+
}
208+
209+
if (*s != 0) {
210+
rt_list_append(res, MP_OBJ_NEW_QSTR(qstr_from_strn_copy(s, strlen(s))));
211+
}
212+
213+
return res;
214+
}
215+
178216
static bool chr_in_str(const char* const str, const size_t str_len, const char c) {
179217
for (size_t i = 0; i < str_len; i++) {
180218
if (str[i] == c) {
@@ -293,12 +331,14 @@ mp_obj_t str_format(uint n_args, const mp_obj_t *args) {
293331

294332
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_find_obj, 2, 4, str_find);
295333
static MP_DEFINE_CONST_FUN_OBJ_2(str_join_obj, str_join);
334+
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_split_obj, 1, 3, str_split);
296335
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_strip_obj, 1, 2, str_strip);
297336
static MP_DEFINE_CONST_FUN_OBJ_VAR(str_format_obj, 1, str_format);
298337

299338
static const mp_method_t str_type_methods[] = {
300339
{ "find", &str_find_obj },
301340
{ "join", &str_join_obj },
341+
{ "split", &str_split_obj },
302342
{ "strip", &str_strip_obj },
303343
{ "format", &str_format_obj },
304344
{ NULL, NULL }, // end-of-list sentinel

tests/basics/string_split.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
print("a b".split())
2+
print(" a b ".split(None))
3+
print(" a b ".split(None, 1))
4+
print(" a b ".split(None, 2))
5+
print(" a b c ".split(None, 1))
6+
print(" a b c ".split(None, 0))
7+
print(" a b c ".split(None, -1))

0 commit comments

Comments
 (0)