py/objstr: Don't treat bytes as unicode in str.count.

jimmo · dpgeorge · commit 9d6f474ea49f · 2022-09-26T00:54:18.000+10:00
`b'\xaa \xaa'.count(b'\xaa')` now (correctly) returns 2 instead of 1. Fixes issue #9404. This work was funded through GitHub Sponsors. Signed-off-by: Jim Mussared <jim.mussared@gmail.com>
diff --git a/py/objstr.c b/py/objstr.c
@@ -1768,14 +1768,16 @@ STATIC mp_obj_t str_count(size_t n_args, const mp_obj_t *args) {
         return MP_OBJ_NEW_SMALL_INT(utf8_charlen(start, end - start) + 1);
     }
 
+    bool is_str = self_type == &mp_type_str;
+
     // count the occurrences
     mp_int_t num_occurrences = 0;
     for (const byte *haystack_ptr = start; haystack_ptr + needle_len <= end;) {
         if (memcmp(haystack_ptr, needle, needle_len) == 0) {
             num_occurrences++;
             haystack_ptr += needle_len;
         } else {
-            haystack_ptr = utf8_next_char(haystack_ptr);
+            haystack_ptr = is_str ? utf8_next_char(haystack_ptr) : haystack_ptr + 1;
         }
     }
 
diff --git a/tests/basics/bytes_count.py b/tests/basics/bytes_count.py
@@ -48,6 +48,13 @@
 print(b"aaaa".count(b'a', -1, 5))
 print(b"abbabba".count(b"abba"))
 
+print(b'\xaa \xaa'.count(b'\xaa'))
+print(b'\xaa \xaa \xaa \xaa'.count(b'\xaa'))
+print(b'\xaa \xaa \xaa \xaa'.count(b'\xaa'), 1)
+print(b'\xaa \xaa \xaa \xaa'.count(b'\xaa'), 2)
+print(b'\xaa \xaa \xaa \xaa'.count(b'\xaa'), 1, 3)
+print(b'\xaa \xaa \xaa \xaa'.count(b'\xaa'), 2, 3)
+
 def t():
     return True
 

Original file line number	Diff line number	Diff line change
`@@ -1768,14 +1768,16 @@ STATIC mp_obj_t str_count(size_t n_args, const mp_obj_t *args) {`
`1768`	`1768`	`return MP_OBJ_NEW_SMALL_INT(utf8_charlen(start, end - start) + 1);`
`1769`	`1769`	`}`
`1770`	`1770`
	`1771`	`+ bool is_str = self_type == &mp_type_str;`
	`1772`	`+`
`1771`	`1773`	`// count the occurrences`
`1772`	`1774`	`mp_int_t num_occurrences = 0;`
`1773`	`1775`	`for (const byte *haystack_ptr = start; haystack_ptr + needle_len <= end;) {`
`1774`	`1776`	`if (memcmp(haystack_ptr, needle, needle_len) == 0) {`
`1775`	`1777`	`num_occurrences++;`
`1776`	`1778`	`haystack_ptr += needle_len;`
`1777`	`1779`	`} else {`
`1778`		`- haystack_ptr = utf8_next_char(haystack_ptr);`
	`1780`	`+ haystack_ptr = is_str ? utf8_next_char(haystack_ptr) : haystack_ptr + 1;`
`1779`	`1781`	`}`
`1780`	`1782`	`}`
`1781`	`1783`