py: In string.count, handle case of zero-length needle.
This commit is contained in:
parent
de4d7aecc8
commit
536dde254b
14
py/objstr.c
14
py/objstr.c
@ -495,8 +495,8 @@ STATIC mp_obj_t str_count(uint n_args, const mp_obj_t *args) {
|
|||||||
GET_STR_DATA_LEN(args[0], haystack, haystack_len);
|
GET_STR_DATA_LEN(args[0], haystack, haystack_len);
|
||||||
GET_STR_DATA_LEN(args[1], needle, needle_len);
|
GET_STR_DATA_LEN(args[1], needle, needle_len);
|
||||||
|
|
||||||
size_t start = 0;
|
machine_uint_t start = 0;
|
||||||
size_t end = haystack_len;
|
machine_uint_t end = haystack_len;
|
||||||
/* TODO use a non-exception-throwing mp_get_index */
|
/* TODO use a non-exception-throwing mp_get_index */
|
||||||
if (n_args >= 3 && args[2] != mp_const_none) {
|
if (n_args >= 3 && args[2] != mp_const_none) {
|
||||||
start = mp_get_index(&str_type, haystack_len, args[2], true);
|
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) {
|
|||||||
end = mp_get_index(&str_type, haystack_len, args[3], true);
|
end = mp_get_index(&str_type, haystack_len, args[3], true);
|
||||||
}
|
}
|
||||||
|
|
||||||
machine_int_t num_occurrences = 0;
|
// if needle_len is zero then we count each gap between characters as an occurrence
|
||||||
|
if (needle_len == 0) {
|
||||||
// needle won't exist in haystack if it's longer, so nothing to count
|
return MP_OBJ_NEW_SMALL_INT(end - start + 1);
|
||||||
if (needle_len > haystack_len) {
|
|
||||||
MP_OBJ_NEW_SMALL_INT(0);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// count the occurrences
|
||||||
|
machine_int_t num_occurrences = 0;
|
||||||
for (machine_uint_t haystack_index = start; haystack_index + needle_len <= end; haystack_index++) {
|
for (machine_uint_t haystack_index = start; haystack_index + needle_len <= end; haystack_index++) {
|
||||||
if (memcmp(&haystack[haystack_index], needle, needle_len) == 0) {
|
if (memcmp(&haystack[haystack_index], needle, needle_len) == 0) {
|
||||||
num_occurrences++;
|
num_occurrences++;
|
||||||
|
@ -1,3 +1,29 @@
|
|||||||
|
print("".count(""))
|
||||||
|
print("".count("a"))
|
||||||
|
print("a".count(""))
|
||||||
|
print("a".count("a"))
|
||||||
|
print("a".count("b"))
|
||||||
|
print("b".count("a"))
|
||||||
|
|
||||||
|
print("aaa".count(""))
|
||||||
|
print("aaa".count("a"))
|
||||||
|
print("aaa".count("aa"))
|
||||||
|
print("aaa".count("aaa"))
|
||||||
|
print("aaa".count("aaaa"))
|
||||||
|
|
||||||
|
print("aaaa".count(""))
|
||||||
|
print("aaaa".count("a"))
|
||||||
|
print("aaaa".count("aa"))
|
||||||
|
print("aaaa".count("aaa"))
|
||||||
|
print("aaaa".count("aaaa"))
|
||||||
|
print("aaaa".count("aaaaa"))
|
||||||
|
|
||||||
|
print("aaa".count("", 1))
|
||||||
|
print("aaa".count("", 2))
|
||||||
|
print("aaa".count("", 3))
|
||||||
|
|
||||||
|
print("aaa".count("", 1, 2))
|
||||||
|
|
||||||
print("asdfasdfaaa".count("asdf", -100))
|
print("asdfasdfaaa".count("asdf", -100))
|
||||||
print("asdfasdfaaa".count("asdf", -8))
|
print("asdfasdfaaa".count("asdf", -8))
|
||||||
print("asdf".count('s', True))
|
print("asdf".count('s', True))
|
||||||
|
Loading…
Reference in New Issue
Block a user