Implement str.count and add tests for it.

Also modify mp_get_index to accept:
1. Indices that are or evaluate to a boolean.
2. Slice indices.
Add tests for these two cases.

1 files changed, 45 insertions, 3 deletions

diff --git a/py/objstr.c b/py/objstr.c
index 130c3266a..c5c7f87f6 100644
--- a/py/objstr.c
+++ b/py/objstr.c
@@ -107,7 +107,7 @@ STATIC mp_obj_t str_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
             // TODO: need predicate to check for int-like type (bools are such for example)
             // ["no", "yes"][1 == 2] is common idiom
             if (MP_OBJ_IS_SMALL_INT(rhs_in)) {
-                uint index = mp_get_index(mp_obj_get_type(lhs_in), lhs_len, rhs_in);
+                uint index = mp_get_index(mp_obj_get_type(lhs_in), lhs_len, rhs_in, false);
                 if (MP_OBJ_IS_TYPE(lhs_in, &bytes_type)) {
                     return MP_OBJ_NEW_SMALL_INT((mp_small_int_t)lhs_data[index]);
                 } else {
@@ -290,10 +290,10 @@ STATIC mp_obj_t str_find(uint n_args, const mp_obj_t *args) {
     size_t end = haystack_len;
     /* TODO use a non-exception-throwing mp_get_index */
     if (n_args >= 3 && args[2] != mp_const_none) {
-        start = mp_get_index(&str_type, haystack_len, args[2]);
+        start = mp_get_index(&str_type, haystack_len, args[2], true);
     }
     if (n_args >= 4 && args[3] != mp_const_none) {
-        end = mp_get_index(&str_type, haystack_len, args[3]);
+        end = mp_get_index(&str_type, haystack_len, args[3], true);
     }
 
     const byte *p = find_subbytes(haystack + start, haystack_len - start, needle, needle_len);
@@ -487,6 +487,46 @@ STATIC mp_obj_t str_replace(uint n_args, const mp_obj_t *args) {
     return mp_obj_str_builder_end(replaced_str);
 }
 
+STATIC mp_obj_t str_count(uint n_args, const mp_obj_t *args) {
+    assert(2 <= n_args && n_args <= 4);
+    assert(MP_OBJ_IS_STR(args[0]));
+    assert(MP_OBJ_IS_STR(args[1]));
+
+    GET_STR_DATA_LEN(args[0], haystack, haystack_len);
+    GET_STR_DATA_LEN(args[1], needle, needle_len);
+
+    size_t start = 0;
+    size_t end = haystack_len;
+    /* TODO use a non-exception-throwing mp_get_index */
+    if (n_args >= 3 && args[2] != mp_const_none) {
+        start = mp_get_index(&str_type, haystack_len, args[2], true);
+    }
+    if (n_args >= 4 && args[3] != mp_const_none) {
+        end = mp_get_index(&str_type, haystack_len, args[3], true);
+    }
+
+    machine_int_t num_occurrences = 0;
+
+    // needle won't exist in haystack if it's longer, so nothing to count
+    if (needle_len > haystack_len) {
+	MP_OBJ_NEW_SMALL_INT(0);
+    }
+
+    for (machine_uint_t haystack_index = start; haystack_index <= end; haystack_index++) {
+	for (machine_uint_t needle_index = 0; needle_index < needle_len; needle_index++) {
+	    if ((haystack_index + needle_len) > end) {
+		return MP_OBJ_NEW_SMALL_INT(num_occurrences);
+	    }
+	    if (haystack[haystack_index + needle_index] == needle[needle_index] && needle_index == (needle_len - 1)) {
+		num_occurrences++;
+	    }
+
+	}
+    }
+
+    return MP_OBJ_NEW_SMALL_INT(num_occurrences);
+}
+
 STATIC machine_int_t str_get_buffer(mp_obj_t self_in, buffer_info_t *bufinfo, int flags) {
     if (flags == BUFFER_READ) {
         GET_STR_DATA_LEN(self_in, str_data, str_len);
@@ -508,6 +548,7 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_startswith_obj, str_startswith);
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_strip_obj, 1, 2, str_strip);
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(str_format_obj, 1, str_format);
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_replace_obj, 3, 4, str_replace);
+STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_count_obj, 2, 4, str_count);
 
 STATIC const mp_method_t str_type_methods[] = {
     { "find", &str_find_obj },
@@ -517,6 +558,7 @@ STATIC const mp_method_t str_type_methods[] = {
     { "strip", &str_strip_obj },
     { "format", &str_format_obj },
     { "replace", &str_replace_obj },
+    { "count", &str_count_obj },
     { NULL, NULL }, // end-of-list sentinel
 };