Skip to content

Commit 2a27365

Browse files
committed
objstr.c: Partial implementation of .rsplit().
sep=None is TODO.
1 parent 51fab28 commit 2a27365

3 files changed

Lines changed: 109 additions & 0 deletions

File tree

py/objstr.c

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
#include "runtime.h"
3939
#include "pfenv.h"
4040
#include "objstr.h"
41+
#include "objlist.h"
4142

4243
STATIC mp_obj_t str_modulo_format(mp_obj_t pattern, uint n_args, const mp_obj_t *args);
4344
const mp_obj_t mp_const_empty_bytes;
@@ -483,6 +484,69 @@ STATIC mp_obj_t str_split(uint n_args, const mp_obj_t *args) {
483484
return res;
484485
}
485486

487+
STATIC mp_obj_t str_rsplit(uint n_args, const mp_obj_t *args) {
488+
if (n_args < 3) {
489+
// If we don't have split limit, it doesn't matter from which side
490+
// we split.
491+
return str_split(n_args, args);
492+
}
493+
const mp_obj_type_t *self_type = mp_obj_get_type(args[0]);
494+
mp_obj_t sep = args[1];
495+
GET_STR_DATA_LEN(args[0], s, len);
496+
497+
machine_int_t splits = mp_obj_get_int(args[2]);
498+
machine_int_t org_splits = splits;
499+
// Preallocate list to the max expected # of elements, as we
500+
// will fill it from the end.
501+
mp_obj_list_t *res = mp_obj_new_list(splits + 1, NULL);
502+
int idx = splits;
503+
504+
if (sep == mp_const_none) {
505+
// TODO
506+
assert(0);
507+
} else {
508+
uint sep_len;
509+
const char *sep_str = mp_obj_str_get_data(sep, &sep_len);
510+
511+
if (sep_len == 0) {
512+
nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "empty separator"));
513+
}
514+
515+
const byte *beg = s;
516+
const byte *last = s + len;
517+
for (;;) {
518+
s = last - sep_len;
519+
for (;;) {
520+
if (splits == 0 || s < beg) {
521+
break;
522+
} else if (memcmp(s, sep_str, sep_len) == 0) {
523+
break;
524+
}
525+
s--;
526+
}
527+
if (s < beg || splits == 0) {
528+
res->items[idx] = str_new(self_type, beg, last - beg);
529+
break;
530+
}
531+
res->items[idx--] = str_new(self_type, s + sep_len, last - s - sep_len);
532+
last = s;
533+
if (splits > 0) {
534+
splits--;
535+
}
536+
}
537+
if (idx != 0) {
538+
// We split less parts than split limit, now go cleanup surplus
539+
int used = org_splits + 1 - idx;
540+
memcpy(res->items, &res->items[idx], used * sizeof(mp_obj_t));
541+
mp_seq_clear(res->items, used, res->alloc, sizeof(*res->items));
542+
res->len = used;
543+
}
544+
}
545+
546+
return res;
547+
}
548+
549+
486550
STATIC mp_obj_t str_finder(uint n_args, const mp_obj_t *args, machine_int_t direction, bool is_index) {
487551
assert(2 <= n_args && n_args <= 4);
488552
assert(MP_OBJ_IS_STR(args[0]));
@@ -1460,6 +1524,7 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_index_obj, 2, 4, str_index);
14601524
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_rindex_obj, 2, 4, str_rindex);
14611525
STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_join_obj, str_join);
14621526
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_split_obj, 1, 3, str_split);
1527+
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_rsplit_obj, 1, 3, str_rsplit);
14631528
STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_startswith_obj, str_startswith);
14641529
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_strip_obj, 1, 2, str_strip);
14651530
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_lstrip_obj, 1, 2, str_lstrip);
@@ -1483,6 +1548,7 @@ STATIC const mp_map_elem_t str_locals_dict_table[] = {
14831548
{ MP_OBJ_NEW_QSTR(MP_QSTR_rindex), (mp_obj_t)&str_rindex_obj },
14841549
{ MP_OBJ_NEW_QSTR(MP_QSTR_join), (mp_obj_t)&str_join_obj },
14851550
{ MP_OBJ_NEW_QSTR(MP_QSTR_split), (mp_obj_t)&str_split_obj },
1551+
{ MP_OBJ_NEW_QSTR(MP_QSTR_rsplit), (mp_obj_t)&str_rsplit_obj },
14861552
{ MP_OBJ_NEW_QSTR(MP_QSTR_startswith), (mp_obj_t)&str_startswith_obj },
14871553
{ MP_OBJ_NEW_QSTR(MP_QSTR_strip), (mp_obj_t)&str_strip_obj },
14881554
{ MP_OBJ_NEW_QSTR(MP_QSTR_lstrip), (mp_obj_t)&str_lstrip_obj },

py/qstrdefs.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -236,6 +236,7 @@ Q(find)
236236
Q(rfind)
237237
Q(rindex)
238238
Q(split)
239+
Q(rsplit)
239240
Q(startswith)
240241
Q(replace)
241242
Q(partition)

tests/basics/string_rsplit.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
# default separator (whitespace)
2+
print("a b".rsplit())
3+
#print(" a b ".rsplit(None))
4+
#print(" a b ".rsplit(None, 1))
5+
#print(" a b ".rsplit(None, 2))
6+
#print(" a b c ".rsplit(None, 1))
7+
#print(" a b c ".rsplit(None, 0))
8+
#print(" a b c ".rsplit(None, -1))
9+
10+
# empty separator should fail
11+
try:
12+
"abc".rsplit('')
13+
except ValueError:
14+
print("ValueError")
15+
16+
# non-empty separator
17+
print("abc".rsplit("a"))
18+
print("abc".rsplit("b"))
19+
print("abc".rsplit("c"))
20+
print("abc".rsplit("z"))
21+
print("abc".rsplit("ab"))
22+
print("abc".rsplit("bc"))
23+
print("abc".rsplit("abc"))
24+
print("abc".rsplit("abcd"))
25+
print("abcabc".rsplit("bc"))
26+
print("abcabc".rsplit("bc", 0))
27+
print("abcabc".rsplit("bc", 1))
28+
print("abcabc".rsplit("bc", 2))
29+
30+
print("10/11/12".rsplit("/", 1))
31+
print("10/11/12".rsplit("/", 2))
32+
print("10/11/12".rsplit("/", 3))
33+
print("10/11/12".rsplit("/", 4))
34+
print("10/11/12".rsplit("/", 5))
35+
36+
print("/*10/*11/*12/*".rsplit("/*", 1))
37+
print("/*10/*11/*12/*".rsplit("/*", 2))
38+
print("/*10/*11/*12/*".rsplit("/*", 3))
39+
print("/*10/*11/*12/*".rsplit("/*", 4))
40+
print("/*10/*11/*12/*".rsplit("/*", 5))
41+
42+
print(b"abcabc".rsplit(b"bc", 2))

0 commit comments

Comments
 (0)