Skip to content

Commit 9e1e8cd

Browse files
xbexbe
committed
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 parent 19438fd commit 9e1e8cd

10 files changed

Lines changed: 119 additions & 22 deletions

File tree

‎py/obj.c‎

Lines changed: 22 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -218,20 +218,32 @@ mp_obj_t *mp_obj_get_array_fixed_n(mp_obj_t o_in, machine_int_t n) {
218218
}
219219
}
220220

221-
uint mp_get_index(const mp_obj_type_t *type, machine_uint_t len, mp_obj_t index) {
222-
// TODO False and True are considered 0 and 1 for indexing purposes
221+
// is_slice determines whether the index is a slice index
222+
uint mp_get_index(const mp_obj_type_t *type, machine_uint_t len, mp_obj_t index, bool is_slice) {
223+
int i;
223224
if (MP_OBJ_IS_SMALL_INT(index)) {
224-
int i = MP_OBJ_SMALL_INT_VALUE(index);
225-
if (i < 0) {
226-
i += len;
227-
}
228-
if (i < 0 || i >= len) {
229-
nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_IndexError, "%s index out of range", qstr_str(type->name)));
230-
}
231-
return i;
225+
i = MP_OBJ_SMALL_INT_VALUE(index);
226+
} else if (MP_OBJ_IS_TYPE(index, &bool_type)) {
227+
i = index == mp_const_true ? 1 : 0;
232228
} else {
233229
nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError, "%s indices must be integers, not %s", qstr_str(type->name), mp_obj_get_type_str(index)));
234230
}
231+
232+
if (i < 0) {
233+
i += len;
234+
}
235+
if (is_slice) {
236+
if (i < 0) {
237+
i = 0;
238+
} else if (i > len) {
239+
i = len;
240+
}
241+
} else {
242+
if (i < 0 || i >= len) {
243+
nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_IndexError, "%s index out of range", qstr_str(type->name)));
244+
}
245+
}
246+
return i;
235247
}
236248

237249
// may return MP_OBJ_NULL

‎py/obj.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -267,7 +267,7 @@ void mp_obj_get_complex(mp_obj_t self_in, mp_float_t *real, mp_float_t *imag);
267267
#endif
268268
//qstr mp_obj_get_qstr(mp_obj_t arg);
269269
mp_obj_t *mp_obj_get_array_fixed_n(mp_obj_t o, machine_int_t n);
270-
uint mp_get_index(const mp_obj_type_t *type, machine_uint_t len, mp_obj_t index);
270+
uint mp_get_index(const mp_obj_type_t *type, machine_uint_t len, mp_obj_t index, bool is_slice);
271271
mp_obj_t mp_obj_len_maybe(mp_obj_t o_in); /* may return NULL */
272272

273273
// none

‎py/objarray.c‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -113,7 +113,7 @@ STATIC mp_obj_t array_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) {
113113
switch (op) {
114114
case RT_BINARY_OP_SUBSCR:
115115
{
116-
uint index = mp_get_index(o->base.type, o->len, rhs);
116+
uint index = mp_get_index(o->base.type, o->len, rhs, false);
117117
return mp_binary_get_val(o->typecode, o->items, index);
118118
}
119119

@@ -140,7 +140,7 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_2(array_append_obj, array_append);
140140

141141
STATIC bool array_store_item(mp_obj_t self_in, mp_obj_t index_in, mp_obj_t value) {
142142
mp_obj_array_t *o = self_in;
143-
uint index = mp_get_index(o->base.type, o->len, index_in);
143+
uint index = mp_get_index(o->base.type, o->len, index_in, false);
144144
mp_binary_set_val(o->typecode, o->items, index, value);
145145
return true;
146146
}

‎py/objlist.c‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ STATIC mp_obj_t list_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) {
104104
return res;
105105
}
106106
#endif
107-
uint index = mp_get_index(o->base.type, o->len, rhs);
107+
uint index = mp_get_index(o->base.type, o->len, rhs, false);
108108
return o->items[index];
109109
}
110110
case RT_BINARY_OP_ADD:
@@ -190,7 +190,7 @@ STATIC mp_obj_t list_pop(uint n_args, const mp_obj_t *args) {
190190
if (self->len == 0) {
191191
nlr_jump(mp_obj_new_exception_msg(&mp_type_IndexError, "pop from empty list"));
192192
}
193-
uint index = mp_get_index(self->base.type, self->len, n_args == 1 ? mp_obj_new_int(-1) : args[1]);
193+
uint index = mp_get_index(self->base.type, self->len, n_args == 1 ? mp_obj_new_int(-1) : args[1], false);
194194
mp_obj_t ret = self->items[index];
195195
self->len -= 1;
196196
memcpy(self->items + index, self->items + index + 1, (self->len - index) * sizeof(mp_obj_t));
@@ -383,7 +383,7 @@ void mp_obj_list_get(mp_obj_t self_in, uint *len, mp_obj_t **items) {
383383

384384
void mp_obj_list_store(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) {
385385
mp_obj_list_t *self = self_in;
386-
uint i = mp_get_index(self->base.type, self->len, index);
386+
uint i = mp_get_index(self->base.type, self->len, index, false);
387387
self->items[i] = value;
388388
}
389389

‎py/objstr.c‎

Lines changed: 45 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@ STATIC mp_obj_t str_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
107107
// TODO: need predicate to check for int-like type (bools are such for example)
108108
// ["no", "yes"][1 == 2] is common idiom
109109
if (MP_OBJ_IS_SMALL_INT(rhs_in)) {
110-
uint index = mp_get_index(mp_obj_get_type(lhs_in), lhs_len, rhs_in);
110+
uint index = mp_get_index(mp_obj_get_type(lhs_in), lhs_len, rhs_in, false);
111111
if (MP_OBJ_IS_TYPE(lhs_in, &bytes_type)) {
112112
return MP_OBJ_NEW_SMALL_INT((mp_small_int_t)lhs_data[index]);
113113
} else {
@@ -290,10 +290,10 @@ STATIC mp_obj_t str_find(uint n_args, const mp_obj_t *args) {
290290
size_t end = haystack_len;
291291
/* TODO use a non-exception-throwing mp_get_index */
292292
if (n_args >= 3 && args[2] != mp_const_none) {
293-
start = mp_get_index(&str_type, haystack_len, args[2]);
293+
start = mp_get_index(&str_type, haystack_len, args[2], true);
294294
}
295295
if (n_args >= 4 && args[3] != mp_const_none) {
296-
end = mp_get_index(&str_type, haystack_len, args[3]);
296+
end = mp_get_index(&str_type, haystack_len, args[3], true);
297297
}
298298

299299
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) {
487487
return mp_obj_str_builder_end(replaced_str);
488488
}
489489

490+
STATIC mp_obj_t str_count(uint n_args, const mp_obj_t *args) {
491+
assert(2 <= n_args && n_args <= 4);
492+
assert(MP_OBJ_IS_STR(args[0]));
493+
assert(MP_OBJ_IS_STR(args[1]));
494+
495+
GET_STR_DATA_LEN(args[0], haystack, haystack_len);
496+
GET_STR_DATA_LEN(args[1], needle, needle_len);
497+
498+
size_t start = 0;
499+
size_t end = haystack_len;
500+
/* TODO use a non-exception-throwing mp_get_index */
501+
if (n_args >= 3 && args[2] != mp_const_none) {
502+
start = mp_get_index(&str_type, haystack_len, args[2], true);
503+
}
504+
if (n_args >= 4 && args[3] != mp_const_none) {
505+
end = mp_get_index(&str_type, haystack_len, args[3], true);
506+
}
507+
508+
machine_int_t num_occurrences = 0;
509+
510+
// needle won't exist in haystack if it's longer, so nothing to count
511+
if (needle_len > haystack_len) {
512+
MP_OBJ_NEW_SMALL_INT(0);
513+
}
514+
515+
for (machine_uint_t haystack_index = start; haystack_index <= end; haystack_index++) {
516+
for (machine_uint_t needle_index = 0; needle_index < needle_len; needle_index++) {
517+
if ((haystack_index + needle_len) > end) {
518+
return MP_OBJ_NEW_SMALL_INT(num_occurrences);
519+
}
520+
if (haystack[haystack_index + needle_index] == needle[needle_index] && needle_index == (needle_len - 1)) {
521+
num_occurrences++;
522+
}
523+
524+
}
525+
}
526+
527+
return MP_OBJ_NEW_SMALL_INT(num_occurrences);
528+
}
529+
490530
STATIC machine_int_t str_get_buffer(mp_obj_t self_in, buffer_info_t *bufinfo, int flags) {
491531
if (flags == BUFFER_READ) {
492532
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);
508548
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_strip_obj, 1, 2, str_strip);
509549
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(str_format_obj, 1, str_format);
510550
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_replace_obj, 3, 4, str_replace);
551+
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_count_obj, 2, 4, str_count);
511552

512553
STATIC const mp_method_t str_type_methods[] = {
513554
{ "find", &str_find_obj },
@@ -517,6 +558,7 @@ STATIC const mp_method_t str_type_methods[] = {
517558
{ "strip", &str_strip_obj },
518559
{ "format", &str_format_obj },
519560
{ "replace", &str_replace_obj },
561+
{ "count", &str_count_obj },
520562
{ NULL, NULL }, // end-of-list sentinel
521563
};
522564

‎py/objtuple.c‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ mp_obj_t tuple_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) {
111111
return res;
112112
}
113113
#endif
114-
uint index = mp_get_index(o->base.type, o->len, rhs);
114+
uint index = mp_get_index(o->base.type, o->len, rhs, false);
115115
return o->items[index];
116116
}
117117
case RT_BINARY_OP_ADD:

‎py/sequence.c‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -149,9 +149,9 @@ mp_obj_t mp_seq_index_obj(const mp_obj_t *items, uint len, uint n_args, const mp
149149
uint stop = len;
150150

151151
if (n_args >= 3) {
152-
start = mp_get_index(type, len, args[2]);
152+
start = mp_get_index(type, len, args[2], true);
153153
if (n_args >= 4) {
154-
stop = mp_get_index(type, len, args[3]);
154+
stop = mp_get_index(type, len, args[3], true);
155155
}
156156
}
157157

‎tests/basics/list_index.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,16 @@
33
print(a.index(2))
44
print(a.index(3))
55
print(a.index(3, 2))
6+
print(a.index(1, -100))
7+
print(a.index(1, False))
8+
9+
try:
10+
print(a.index(1, True))
11+
except ValueError:
12+
print("Raised ValueError")
13+
else:
14+
print("Did not raise ValueError")
15+
616
try:
717
print(a.index(3, 2, 2))
818
except ValueError:

‎tests/basics/string_count.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
print("asdfasdfaaa".count("asdf", -100))
2+
print("asdfasdfaaa".count("asdf", -8))
3+
print("asdf".count('s', True))
4+
print("asdf".count('a', True))
5+
print("asdf".count('a', False))
6+
print("asdf".count('a', 1 == 2))
7+
print("hello world".count('l'))
8+
print("hello world".count('l', 5))
9+
print("hello world".count('l', 3))
10+
print("hello world".count('z', 3, 6))
11+
print("aaaa".count('a'))
12+
print("aaaa".count('a', 0, 3))
13+
print("aaaa".count('a', 0, 4))
14+
print("aaaa".count('a', 0, 5))
15+
print("aaaa".count('a', 1, 5))
16+
print("aaaa".count('a', -1, 5))
17+
18+
def t():
19+
return True
20+
21+
print("0000".count('0', t()))

‎tests/basics/string_find.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,3 +9,15 @@
99
print("hello world".find("ll", 1, 3))
1010
print("hello world".find("ll", 1, 4))
1111
print("hello world".find("ll", 1, 5))
12+
print("hello world".find("ll", -100))
13+
print("0000".find('0'))
14+
print("0000".find('0', 0))
15+
print("0000".find('0', 1))
16+
print("0000".find('0', 2))
17+
print("0000".find('0', 3))
18+
print("0000".find('0', 4))
19+
print("0000".find('0', 5))
20+
print("0000".find('-1', 3))
21+
print("0000".find('1', 3))
22+
print("0000".find('1', 4))
23+
print("0000".find('1', 5))

0 commit comments

Comments
 (0)