diff --git a/py/obj.h b/py/obj.h index 965155553f..48d332e200 100644 --- a/py/obj.h +++ b/py/obj.h @@ -403,3 +403,4 @@ bool m_seq_get_fast_slice_indexes(machine_uint_t len, mp_obj_t slice, machine_ui bool mp_seq_cmp_bytes(int op, const byte *data1, uint len1, const byte *data2, uint len2); bool mp_seq_cmp_objs(int op, const mp_obj_t *items1, uint len1, const mp_obj_t *items2, uint len2); mp_obj_t mp_seq_index_obj(const mp_obj_t *items, uint len, uint n_args, const mp_obj_t *args); +mp_obj_t mp_seq_count_obj(const mp_obj_t *items, uint len, mp_obj_t value); diff --git a/py/objlist.c b/py/objlist.c index 0c6525be8a..4e8944d35e 100644 --- a/py/objlist.c +++ b/py/objlist.c @@ -260,14 +260,7 @@ static mp_obj_t list_copy(mp_obj_t self_in) { static mp_obj_t list_count(mp_obj_t self_in, mp_obj_t value) { assert(MP_OBJ_IS_TYPE(self_in, &list_type)); mp_obj_list_t *self = self_in; - int count = 0; - for (int i = 0; i < self->len; i++) { - if (mp_obj_equal(self->items[i], value)) { - count++; - } - } - - return mp_obj_new_int(count); + return mp_seq_count_obj(self->items, self->len, value); } static mp_obj_t list_index(uint n_args, const mp_obj_t *args) { diff --git a/py/objtuple.c b/py/objtuple.c index 581bfb6290..7221db7746 100644 --- a/py/objtuple.c +++ b/py/objtuple.c @@ -153,6 +153,13 @@ static mp_obj_t tuple_getiter(mp_obj_t o_in) { return mp_obj_new_tuple_iterator(o_in, 0); } +static mp_obj_t tuple_count(mp_obj_t self_in, mp_obj_t value) { + assert(MP_OBJ_IS_TYPE(self_in, &tuple_type)); + mp_obj_tuple_t *self = self_in; + return mp_seq_count_obj(self->items, self->len, value); +} +static MP_DEFINE_CONST_FUN_OBJ_2(tuple_count_obj, tuple_count); + static mp_obj_t tuple_index(uint n_args, const mp_obj_t *args) { assert(MP_OBJ_IS_TYPE(args[0], &tuple_type)); mp_obj_tuple_t *self = args[0]; @@ -161,6 +168,7 @@ static mp_obj_t tuple_index(uint n_args, const mp_obj_t *args) { static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(tuple_index_obj, 2, 4, tuple_index); static const mp_method_t tuple_type_methods[] = { + { "count", &tuple_count_obj }, { "index", &tuple_index_obj }, { NULL, NULL }, // end-of-list sentinel }; diff --git a/py/sequence.c b/py/sequence.c index b3fb9fcd34..d8cfd9f3e9 100644 --- a/py/sequence.c +++ b/py/sequence.c @@ -164,3 +164,15 @@ mp_obj_t mp_seq_index_obj(const mp_obj_t *items, uint len, uint n_args, const mp nlr_jump(mp_obj_new_exception_msg(MP_QSTR_ValueError, "object not in sequence")); } + +mp_obj_t mp_seq_count_obj(const mp_obj_t *items, uint len, mp_obj_t value) { + uint count = 0; + for (uint i = 0; i < len; i++) { + if (mp_obj_equal(items[i], value)) { + count++; + } + } + + // Common sense says this cannot overflow small int + return MP_OBJ_NEW_SMALL_INT(count); +} diff --git a/tests/basics/tuple_count.py b/tests/basics/tuple_count.py new file mode 100644 index 0000000000..7f42ede28c --- /dev/null +++ b/tests/basics/tuple_count.py @@ -0,0 +1,5 @@ +a = (1, 2, 3) +a = a + a + a +b = (0, 0, a, 0, a, 0) +print(a.count(2)) +print(b.count(a))