py/objstr: Support tuples and start/end args in startswith and endswith.

This change allows tuples to be passed as the prefix/suffix argument to the
`str.startswith()` and `str.endswith()` methods.  The methods will return
`True` if the string starts/ends with any of the prefixes/suffixes in the
tuple.

Also adds full support for the `start` and `end` arguments to both methods
for compatibility with CPython.

Tests have been updated for the new behaviour.

Signed-off-by: Glenn Moloney <[email protected]>

Author: glenn20

py/objstr.c

@@ -67,6 +67,26 @@ static void check_is_str_or_bytes(mp_obj_t self_in) {
     mp_check_self(mp_obj_is_str_or_bytes(self_in));
 }
 
+static const byte *get_substring_data(const mp_obj_t obj, size_t n_args, const mp_obj_t *args, size_t *len) {
+    // Get substring data from obj, using args[0,1] to specify start and end indices.
+    GET_STR_DATA_LEN(obj, str, str_len);
+    if (n_args > 0) {
+        const mp_obj_type_t *self_type = mp_obj_get_type(obj);
+        const byte *end = str + str_len;
+        if (n_args > 1 && args[1] != mp_const_none) {
+            end = str_index_to_ptr(self_type, str, str_len, args[1], true);
+        }
+        if (args[0] != mp_const_none) {
+            str = str_index_to_ptr(self_type, str, str_len, args[0], true);
+        }
+        str_len = MAX(end - str, 0);
+    }
+    if (len) {
+        *len = str_len;
+    }
+    return str;
+}
+
 /******************************************************************************/
 /* str                                                                        */
 
@@ -802,37 +822,34 @@ static mp_obj_t str_rindex(size_t n_args, const mp_obj_t *args) {
 }
 MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_rindex_obj, 2, 4, str_rindex);
 
-// TODO: (Much) more variety in args
-static mp_obj_t str_startswith(size_t n_args, const mp_obj_t *args) {
-    const mp_obj_type_t *self_type = mp_obj_get_type(args[0]);
-    GET_STR_DATA_LEN(args[0], str, str_len);
-    size_t prefix_len;
-    const char *prefix = mp_obj_str_get_data(args[1], &prefix_len);
-    const byte *start = str;
-    if (n_args > 2) {
-        start = str_index_to_ptr(self_type, str, str_len, args[2], true);
+static mp_obj_t str_startendswith(size_t n_args, const mp_obj_t *args, bool ends_with) {
+    size_t str_len;
+    const byte *str = get_substring_data(args[0], n_args - 2, args + 2, &str_len);
+    mp_obj_t *prefixes = (mp_obj_t *)&args[1];
+    size_t n_prefixes = 1;
+    if (mp_obj_is_type(args[1], &mp_type_tuple)) {
+        mp_obj_tuple_get(args[1], &n_prefixes, &prefixes);
     }
-    if (prefix_len + (start - str) > str_len) {
-        return mp_const_false;
+    size_t prefix_len;
+    for (size_t i = 0; i < n_prefixes; i++) {
+        const char *prefix = mp_obj_str_get_data(prefixes[i], &prefix_len);
+        const byte *s = str + (ends_with ? str_len - prefix_len : 0);
+        if (prefix_len <= str_len && memcmp(s, prefix, prefix_len) == 0) {
+            return mp_const_true;
+        }
     }
-    return mp_obj_new_bool(memcmp(start, prefix, prefix_len) == 0);
+    return mp_const_false;
 }
-MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_startswith_obj, 2, 3, str_startswith);
 
-static mp_obj_t str_endswith(size_t n_args, const mp_obj_t *args) {
-    GET_STR_DATA_LEN(args[0], str, str_len);
-    size_t suffix_len;
-    const char *suffix = mp_obj_str_get_data(args[1], &suffix_len);
-    if (n_args > 2) {
-        mp_raise_NotImplementedError(MP_ERROR_TEXT("start/end indices"));
-    }
+static mp_obj_t str_startswith(size_t n_args, const mp_obj_t *args) {
+    return str_startendswith(n_args, args, false);
+}
+MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_startswith_obj, 2, 4, str_startswith);
 
-    if (suffix_len > str_len) {
-        return mp_const_false;
-    }
-    return mp_obj_new_bool(memcmp(str + (str_len - suffix_len), suffix, suffix_len) == 0);
+static mp_obj_t str_endswith(size_t n_args, const mp_obj_t *args) {
+    return str_startendswith(n_args, args, true);
 }
-MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_endswith_obj, 2, 3, str_endswith);
+MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_endswith_obj, 2, 4, str_endswith);
 
 enum { LSTRIP, RSTRIP, STRIP };
 

tests/basics/string_endswith.py

@@ -5,11 +5,28 @@
 print("foobar".endswith(""))
 print("foobar".endswith("foobarbaz"))
 
-#print("1foobar".startswith("foo", 1))
-#print("1foo".startswith("foo", 1))
-#print("1foo".startswith("1foo", 1))
-#print("1fo".startswith("foo", 1))
-#print("1fo".startswith("foo", 10))
+print("foobar".endswith("bar", 3))
+print("foobar".endswith("bar", 4))
+print("foobar".endswith("foo", 0, 3))
+print("foobar".endswith("foo", 0, 4))
+print("foobar".endswith("foo", 1, 3))
+print("foobar".endswith("foo", 1, 3))
+print("foobar".endswith("oo", 1, 3))
+print("foobar".endswith("o", 2, 3))
+print("foobar".endswith("o", 3, 3))
+print("foobar".endswith("o", 4, 3))
+
+print("foobar".endswith("bar", None, None))
+print("foobar".endswith("bar", None, 3))
+print("foobar".endswith("bar", 3, None))
+print("foobar".endswith("bar", 2, None))
+print("foobar".endswith("foo", None, 3))
+
+print("foobar".endswith(("bar", "foo")))
+print("foobar".endswith(("foo", "bar")))
+print("foobar".endswith(("foo", "bar1")))
+print("foobar".endswith(("bar", )))
+print("foobar".endswith(("foo", )))
 
 try:
     "foobar".endswith(1)

tests/basics/string_endswith_upy.py

@@ -10,6 +10,25 @@
 print("1fo".startswith("foo", 1))
 print("1fo".startswith("foo", 10))
 
+print("1foobar".startswith("foo", 1, 5))
+print("1foobar".startswith("foo", 1, 4))
+print("1foobar".startswith("foo", 1, 3))
+print("1foobar".startswith("oo", 2, 4))
+print("1foobar".startswith("o", 3, 4))
+print("1foobar".startswith("o", 4, 4))
+print("1foobar".startswith("o", 5, 4))
+
+print("foobar".startswith("foo", None, None))
+print("foobar".startswith("foo", None, 3))
+print("foobar".startswith("foo", None, 2))
+print("foobar".startswith("bar", 3, None))
+
+
+print("foobar".startswith(("foo", "sth")))
+print("foobar".startswith(("sth", "foo")))
+print("foobar".startswith(("sth", "foo2")))
+print("foobar".startswith(("foo", )))
+
 try:
     "foobar".startswith(1)
 except TypeError:

tests/basics/string_endswith_upy.py.exp

@@ -63,12 +63,6 @@
 except NotImplementedError:
     print("NotImplementedError")
 
-# str.endswith(s, start) not implemented
-try:
-    "abc".endswith("c", 1)
-except NotImplementedError:
-    print("NotImplementedError")
-
 # str subscr with step!=1 not implemented
 try:
     print("abc"[1:2:3])

tests/basics/string_startswith.py

@@ -13,7 +13,6 @@ NotImplementedError
 NotImplementedError
 NotImplementedError
 NotImplementedError
-NotImplementedError
 b'\x01\x02'
 b'\x01\x00'
 NotImplementedError