8000 py/objstr: Support tuples and start/end args in startswith and endswith. · sparkfun/micropython@eb45d97 · GitHub
[go: up one dir, main page]
Skip to content

Commit eb45d97

Browse files
glenn20dpgeorge
glenn20
authored and
dpgeorge
committed
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 <glenn.moloney@gmail.com>
1 parent 69ffd2a commit eb45d97

9 files changed

+83
-51
lines changed

‎py/objstr.c

Lines changed: 42 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,26 @@ static void check_is_str_or_bytes(mp_obj_t self_in) {
6767
mp_check_self(mp_obj_is_str_or_bytes(self_in));
6868
}
6969

70+
static const byte *get_substring_data(const mp_obj_t obj, size_t n_args, const mp_obj_t *args, size_t *len) {
71+
// Get substring data from obj, using args[0,1] to specify start and end indices.
72+
GET_STR_DATA_LEN(obj, str, str_len);
73+
if (n_args > 0) {
74+
const mp_obj_type_t *self_type = mp_obj_get_type(obj);
75+
const byte *end = str + str_len;
76+
if (n_args > 1 && args[1] != mp_const_none) {
77+
end = str_index_to_ptr(self_type, str, str_len, args[1], true);
78+
}
79+
if (args[0] != mp_const_none) {
80+
str = str_index_to_ptr(self_type, str, str_len, args[0], true);
81+
}
82+
str_len = MAX(end - str, 0);
83+
}
84+
if (len) {
85+
*len = str_len;
86+
}
87+
return str;
88+
}
89+
7090
/******************************************************************************/
7191
/* str */
7292

@@ -802,37 +822,34 @@ static mp_obj_t str_rindex(size_t n_args, const mp_obj_t *args) {
802822
}
803823
MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_rindex_obj, 2, 4, str_rindex);
804824

805-
// TODO: (Much) more variety in args
806-
static mp_obj_t str_startswith(size_t n_args, const mp_obj_t *args) {
807-
const mp_obj_type_t *self_type = mp_obj_get_type(args[0]);
808-
GET_STR_DATA_LEN(args[0], str, str_len);
809-
size_t prefix_len;
810-
const char *prefix = mp_obj_str_get_data(args[1], &prefix_len);
811-
const byte *start = str;
812-
if (n_args > 2) {
813-
start = str_index_to_ptr(self_type, str, str_len, args[2], true);
825+
static mp_obj_t str_startendswith(size_t n_args, const mp_obj_t *args, bool ends_with) {
826+
size_t str_len;
827+
const byte *str = get_substring_data(args[0], n_args - 2, args + 2, &str_len);
828+
mp_obj_t *prefixes = (mp_obj_t *)&args[1];
829+
size_t n_prefixes = 1;
830+
if (mp_obj_is_type(args[1], &mp_type_tuple)) {
831+
mp_obj_tuple_get(args[1], &n_prefixes, &prefixes);
814832
}
815-
if (prefix_len + (start - str) > str_len) {
816-
return mp_const_false;
833+
size_t prefix_len;
834+
for (size_t i = 0; i < n_prefixes; i++) {
835+
const char *prefix = mp_obj_str_get_data(prefixes[i], &prefix_len);
836+
const byte *s = str + (ends_with ? str_len - prefix_len : 0);
837+
if (prefix_len <= str_len && memcmp(s, prefix, prefix_len) == 0) {
838+
return mp_const_true;
839+
}
817840
}
818-
return mp_obj_new_bool(memcmp(start, prefix, prefix_len) == 0);
841+
return mp_const_false;
819842
}
820-
MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_startswith_obj, 2, 3, str_startswith);
821843

822-
static mp_obj_t str_endswith(size_t n_args, const mp_obj_t *args) {
823-
GET_STR_DATA_LEN(args[0], str, str_len);
824-
size_t suffix_len;
825-
const char *suffix = mp_obj_str_get_data(args[1], &suffix_len);
826-
if (n_args > 2) {
827-
mp_raise_NotImplementedError(MP_ERROR_TEXT("start/end indices"));
828-
}
844+
static mp_obj_t str_startswith(size_t n_args, const mp_obj_t *args) {
845+
return str_startendswith(n_args, args, false);
846+
}
847+
MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_startswith_obj, 2, 4, str_startswith);
829848

830-
if (suffix_len > str_len) {
831-
return mp_const_false;
832-
}
833-
return mp_obj_new_bool(memcmp(str + (str_len - suffix_len), suffix, suffix_len) == 0);
849+
static mp_obj_t str_endswith(size_t n_args, const mp_obj_t *args) {
850+
return str_startendswith(n_args, args, true);
834851
}
835-
MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_endswith_obj, 2, 3, str_endswith);
852+
MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_endswith_obj, 2, 4, str_endswith);
836853

837854
enum { LSTRIP, RSTRIP, STRIP };
838855

‎tests/basics/string_endswith.py

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,28 @@
55
print("foobar".endswith(""))
66
print("foobar".endswith("foobarbaz"))
77

8-
#print("1foobar".startswith("foo", 1))
9-
#print("1foo".startswith("foo", 1))
10-
#print("1foo".startswith("1foo", 1))
11-
#print("1fo".startswith("foo", 1))
12-
#print("1fo".startswith("foo", 10))
8+
print("foobar".endswith("bar", 3))
9+
print("foobar".endswith("bar", 4))
10+
print("foobar".endswith("foo", 0, 3))
11+
print("foobar".endswith("foo", 0, 4))
12+
print("foobar".endswith("foo", 1, 3))
13+
print("foobar".endswith("foo", 1, 3))
14+
print("foobar".endswith("oo", 1, 3))
15+
print("foobar".endswith("o", 2, 3))
16+
print("foobar".endswith("o", 3, 3))
17+
print("foobar".endswith("o", 4, 3))
18+
19+
print("foobar".endswith("bar", None, None))
20+
print("foobar".endswith("bar", None, 3))
21+
print("foobar".endswith("bar", 3, None))
22+
print("foobar".endswith("bar", 2, None))
23+
print("foobar".endswith("foo", None, 3))
24+
25+
print("foobar".endswith(("bar", "foo")))
26+
print("foobar".endswith(("foo", "bar")))
27+
print("foobar".endswith(("foo", "bar1")))
28+
print("foobar".endswith(("bar", )))
29+
print("foobar".endswith(("foo", )))
1330

1431
try:
1532
"foobar".endswith(1)

‎tests/basics/string_endswith_upy.py

Lines changed: 0 additions & 6 deletions
This file was deleted.

‎tests/basics/string_endswith_upy.py.exp

Lines changed: 0 additions & 1 deletion
This file was deleted.

‎tests/basics/string_startswith.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,25 @@
1010
print("1fo".startswith("foo", 1))
1111
print("1fo".startswith("foo", 10))
1212

13+
print("1foobar".startswith("foo", 1, 5))
14+
print("1foobar".startswith("foo", 1, 4))
15+
print("1foobar".startswith("foo", 1, 3))
16+
print("1foobar".startswith("oo", 2, 4))
17+
print("1foobar".startswith("o", 3, 4))
18+
print("1foobar".startswith("o", 4, 4))
19+
print("1foobar".startswith("o", 5, 4))
20+
21+
print("foobar".startswith("foo", None, None))
22+
print("foobar" 83CB .startswith("foo", None, 3))
23+
print("foobar".startswith("foo", None, 2))
24+
print("foobar".startswith("bar", 3, None))
25+
26+
27+
print("foobar".startswith(("foo", "sth")))
28+
print("foobar".startswith(("sth", "foo")))
29+
print("foobar".startswith(("sth", "foo2")))
30+
print("foobar".startswith(("foo", )))
31+
1332
try:
1433
"foobar".startswith(1)
1534
except TypeError:

‎tests/basics/string_startswith_upy.py

Lines changed: 0 additions & 6 deletions
This file was deleted.

‎tests/basics/string_startswith_upy.py.exp

Lines changed: 0 additions & 1 deletion
This file was deleted.

‎tests/misc/non_compliant.py

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -63,12 +63,6 @@
6363
except NotImplementedError:
6464
print("NotImplementedError")
6565

66-
# str.endswith(s, start) not implemented
67-
try:
68-
"abc".endswith("c", 1)
69-
except NotImplementedError:
70-
print("NotImplementedError")
71-
7266
# str subscr with step!=1 not implemented
7367
try:
7468
print("abc"[1:2:3])

‎tests/misc/non_compliant.py.exp

Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@ NotImplementedError
1313
NotImplementedError
1414
NotImplementedError
1515
NotImplementedError
16-
NotImplementedError
1716
b'\x01\x02'
1817
b'\x01\x00'
1918
NotImplementedError

0 commit comments

Comments
 (0)
0