Skip to content

Commit 536dde2

Browse files
dpgeorgedpgeorge
committed
py: In string.count, handle case of zero-length needle.
1 parent de4d7ae commit 536dde2

2 files changed

Lines changed: 33 additions & 7 deletions

File tree

‎py/objstr.c‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -495,8 +495,8 @@ STATIC mp_obj_t str_count(uint n_args, const mp_obj_t *args) {
495495
GET_STR_DATA_LEN(args[0], haystack, haystack_len);
496496
GET_STR_DATA_LEN(args[1], needle, needle_len);
497497

498-
size_t start = 0;
499-
size_t end = haystack_len;
498+
machine_uint_t start = 0;
499+
machine_uint_t end = haystack_len;
500500
/* TODO use a non-exception-throwing mp_get_index */
501501
if (n_args >= 3 && args[2] != mp_const_none) {
502502
start = mp_get_index(&str_type, haystack_len, args[2], true);
@@ -505,13 +505,13 @@ STATIC mp_obj_t str_count(uint n_args, const mp_obj_t *args) {
505505
end = mp_get_index(&str_type, haystack_len, args[3], true);
506506
}
507507

508-
machine_int_t num_occurrences = 0;
509-
510-
// needle won't exist in haystack if it's longer, so nothing to count
511-
if (needle_len > haystack_len) {
512-
MP_OBJ_NEW_SMALL_INT(0);
508+
// if needle_len is zero then we count each gap between characters as an occurrence
509+
if (needle_len == 0) {
510+
return MP_OBJ_NEW_SMALL_INT(end - start + 1);
513511
}
514512

513+
// count the occurrences
514+
machine_int_t num_occurrences = 0;
515515
for (machine_uint_t haystack_index = start; haystack_index + needle_len <= end; haystack_index++) {
516516
if (memcmp(&haystack[haystack_index], needle, needle_len) == 0) {
517517
num_occurrences++;

‎tests/basics/string_count.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,29 @@
1+
print("".count(""))
2+
print("".count("a"))
3+
print("a".count(""))
4+
print("a".count("a"))
5+
print("a".count("b"))
6+
print("b".count("a"))
7+
8+
print("aaa".count(""))
9+
print("aaa".count("a"))
10+
print("aaa".count("aa"))
11+
print("aaa".count("aaa"))
12+
print("aaa".count("aaaa"))
13+
14+
print("aaaa".count(""))
15+
print("aaaa".count("a"))
16+
print("aaaa".count("aa"))
17+
print("aaaa".count("aaa"))
18+
print("aaaa".count("aaaa"))
19+
print("aaaa".count("aaaaa"))
20+
21+
print("aaa".count("", 1))
22+
print("aaa".count("", 2))
23+
print("aaa".count("", 3))
24+
25+
print("aaa".count("", 1, 2))
26+
127
print("asdfasdfaaa".count("asdf", -100))
228
print("asdfasdfaaa".count("asdf", -8))
329
print("asdf".count('s', True))

0 commit comments

Comments
 (0)