Skip to content

Commit e3e7c2b

Browse files
dpgeorgedpgeorge
committed
Merge pull request adafruit#351 from xbe/str-partition
Implement str.partition and add tests for it.
2 parents dfbafab + 0a6894c commit e3e7c2b

3 files changed

Lines changed: 102 additions & 0 deletions

File tree

‎py/objstr.c‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -520,6 +520,46 @@ STATIC mp_obj_t str_count(uint n_args, const mp_obj_t *args) {
520520
return MP_OBJ_NEW_SMALL_INT(num_occurrences);
521521
}
522522

523+
STATIC mp_obj_t str_partitioner(mp_obj_t self_in, mp_obj_t arg, bool rpartition) {
524+
assert(MP_OBJ_IS_STR(self_in));
525+
if (!MP_OBJ_IS_STR(arg)) {
526+
nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError,
527+
"Can't convert '%s' object to str implicitly", mp_obj_get_type_str(arg)));
528+
}
529+
GET_STR_DATA_LEN(self_in, str, str_len);
530+
GET_STR_DATA_LEN(arg, sep, sep_len);
531+
mp_obj_t result[] = {MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_)};
532+
533+
if (sep_len == 0) {
534+
nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "empty separator"));
535+
}
536+
if (rpartition) {
537+
result[2] = mp_obj_new_str(str, str_len, false);
538+
} else {
539+
result[0] = mp_obj_new_str(str, str_len, false);
540+
}
541+
542+
for (machine_uint_t str_index = 0; str_index + sep_len <= str_len; str_index++) {
543+
if (memcmp(&str[str_index], sep, sep_len) == 0) {
544+
result[0] = mp_obj_new_str(str, str_index, false);
545+
result[1] = arg;
546+
result[2] = mp_obj_new_str(str + str_index + sep_len, str_len - str_index - sep_len, false);
547+
if (!rpartition) {
548+
break;
549+
}
550+
}
551+
}
552+
return mp_obj_new_tuple(3, result);
553+
}
554+
555+
STATIC mp_obj_t str_partition(mp_obj_t self_in, mp_obj_t arg, bool partition) {
556+
return str_partitioner(self_in, arg, false);
557+
}
558+
559+
STATIC mp_obj_t str_rpartition(mp_obj_t self_in, mp_obj_t arg, bool partition) {
560+
return str_partitioner(self_in, arg, true);
561+
}
562+
523563
STATIC machine_int_t str_get_buffer(mp_obj_t self_in, buffer_info_t *bufinfo, int flags) {
524564
if (flags == BUFFER_READ) {
525565
GET_STR_DATA_LEN(self_in, str_data, str_len);
@@ -542,6 +582,8 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_strip_obj, 1, 2, str_strip);
542582
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(str_format_obj, 1, str_format);
543583
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_replace_obj, 3, 4, str_replace);
544584
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_count_obj, 2, 4, str_count);
585+
STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_partition_obj, str_partition);
586+
STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_rpartition_obj, str_rpartition);
545587

546588
STATIC const mp_method_t str_type_methods[] = {
547589
{ "find", &str_find_obj },
@@ -552,6 +594,8 @@ STATIC const mp_method_t str_type_methods[] = {
552594
{ "format", &str_format_obj },
553595
{ "replace", &str_replace_obj },
554596
{ "count", &str_count_obj },
597+
{ "partition", &str_partition_obj },
598+
{ "rpartition", &str_rpartition_obj },
555599
{ NULL, NULL }, // end-of-list sentinel
556600
};
557601

‎tests/basics/string_partition.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
print("asdf".partition('g'))
2+
print("asdf".partition('a'))
3+
print("asdf".partition('s'))
4+
print("asdf".partition('f'))
5+
print("asdf".partition('d'))
6+
print("asdf".partition('asd'))
7+
print("asdf".partition('sdf'))
8+
print("asdf".partition('as'))
9+
print("asdf".partition('df'))
10+
print("asdf".partition('asdf'))
11+
print("asdf".partition('asdfa'))
12+
print("asdf".partition('fasdf'))
13+
print("asdf".partition('fasdfa'))
14+
print("abba".partition('a'))
15+
print("abba".partition('b'))
16+
17+
try:
18+
print("asdf".partition(1))
19+
except TypeError:
20+
print("Raised TypeError")
21+
else:
22+
print("Did not raise TypeError")
23+
24+
try:
25+
print("asdf".partition(''))
26+
except ValueError:
27+
print("Raised ValueError")
28+
else:
29+
print("Did not raise ValueError")

‎tests/basics/string_rpartition.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
print("asdf".rpartition('g'))
2+
print("asdf".rpartition('a'))
3+
print("asdf".rpartition('s'))
4+
print("asdf".rpartition('f'))
5+
print("asdf".rpartition('d'))
6+
print("asdf".rpartition('asd'))
7+
print("asdf".rpartition('sdf'))
8+
print("asdf".rpartition('as'))
9+
print("asdf".rpartition('df'))
10+
print("asdf".rpartition('asdf'))
11+
print("asdf".rpartition('asdfa'))
12+
print("asdf".rpartition('fasdf'))
13+
print("asdf".rpartition('fasdfa'))
14+
print("abba".rpartition('a'))
15+
print("abba".rpartition('b'))
16+
17+
try:
18+
print("asdf".rpartition(1))
19+
except TypeError:
20+
print("Raised TypeError")
21+
else:
22+
print("Did not raise TypeError")
23+
24+
try:
25+
print("asdf".rpartition(''))
26+
except ValueError:
27+
print("Raised ValueError")
28+
else:
29+
print("Did not raise ValueError")

0 commit comments

Comments
 (0)