diff --git a/py/obj.c b/py/obj.c index fb59eec82c..322a302f97 100644 --- a/py/obj.c +++ b/py/obj.c @@ -531,6 +531,36 @@ mp_obj_t mp_identity_getiter(mp_obj_t self, mp_obj_iter_buf_t *iter_buf) { return self; } +typedef struct { + mp_obj_base_t base; + mp_fun_1_t iternext; + mp_obj_t obj; + mp_int_t cur; +} mp_obj_generic_it_t; + +STATIC mp_obj_t generic_it_iternext(mp_obj_t self_in) { + mp_obj_generic_it_t *self = MP_OBJ_TO_PTR(self_in); + mp_obj_type_t *type = mp_obj_get_type(self->obj); + mp_obj_t current_length = type->unary_op(MP_UNARY_OP_LEN, self->obj); + if (self->cur < MP_OBJ_SMALL_INT_VALUE(current_length)) { + mp_obj_t o_out = type->subscr(self->obj, MP_OBJ_NEW_SMALL_INT(self->cur), MP_OBJ_SENTINEL); + self->cur += 1; + return o_out; + } else { + return MP_OBJ_STOP_ITERATION; + } +} + +mp_obj_t mp_obj_new_generic_iterator(mp_obj_t obj, mp_obj_iter_buf_t *iter_buf) { + assert(sizeof(mp_obj_generic_it_t) <= sizeof(mp_obj_iter_buf_t)); + mp_obj_generic_it_t *o = (mp_obj_generic_it_t*)iter_buf; + o->base.type = &mp_type_polymorph_iter; + o->iternext = generic_it_iternext; + o->obj = obj; + o->cur = 0; + return MP_OBJ_FROM_PTR(o); +} + bool mp_get_buffer(mp_obj_t obj, mp_buffer_info_t *bufinfo, mp_uint_t flags) { mp_obj_type_t *type = mp_obj_get_type(obj); if (type->buffer_p.get_buffer == NULL) { diff --git a/py/obj.h b/py/obj.h index ae2bff1906..773043b593 100644 --- a/py/obj.h +++ b/py/obj.h @@ -819,6 +819,10 @@ mp_obj_t mp_identity(mp_obj_t self); MP_DECLARE_CONST_FUN_OBJ_1(mp_identity_obj); mp_obj_t mp_identity_getiter(mp_obj_t self, mp_obj_iter_buf_t *iter_buf); +// Generic iterator that uses unary op and subscr to iterate over a native type. It will be slower +// than a custom iterator but applies broadly. +mp_obj_t mp_obj_new_generic_iterator(mp_obj_t self, mp_obj_iter_buf_t *iter_buf); + // module typedef struct _mp_obj_module_t { mp_obj_base_t base; diff --git a/shared-bindings/displayio/Group.c b/shared-bindings/displayio/Group.c index bbf7190a3f..76719c5808 100644 --- a/shared-bindings/displayio/Group.c +++ b/shared-bindings/displayio/Group.c @@ -195,6 +195,21 @@ STATIC mp_obj_t displayio_group_obj_insert(mp_obj_t self_in, mp_obj_t index_obj, } MP_DEFINE_CONST_FUN_OBJ_3(displayio_group_insert_obj, displayio_group_obj_insert); + +//| .. method:: index(layer) +//| +//| Returns the index of the first copy of layer. Raises ValueError if not found. +//| +STATIC mp_obj_t displayio_group_obj_index(mp_obj_t self_in, mp_obj_t layer) { + displayio_group_t *self = native_group(self_in); + mp_int_t index = common_hal_displayio_group_index(self, layer); + if (index < 0) { + mp_raise_ValueError(translate("object not in sequence")); + } + return MP_OBJ_NEW_SMALL_INT(index); +} +MP_DEFINE_CONST_FUN_OBJ_2(displayio_group_index_obj, displayio_group_obj_index); + //| .. method:: pop(i=-1) //| //| Remove the ith item and return it. @@ -217,6 +232,20 @@ STATIC mp_obj_t displayio_group_obj_pop(size_t n_args, const mp_obj_t *pos_args, } MP_DEFINE_CONST_FUN_OBJ_KW(displayio_group_pop_obj, 1, displayio_group_obj_pop); + +//| .. method:: remove(layer) +//| +//| Remove the first copy of layer. Raises ValueError if it is not present. +//| +STATIC mp_obj_t displayio_group_obj_remove(mp_obj_t self_in, mp_obj_t layer) { + mp_obj_t index = displayio_group_obj_index(self_in, layer); + displayio_group_t *self = native_group(self_in); + + common_hal_displayio_group_pop(self, MP_OBJ_SMALL_INT_VALUE(index)); + return mp_const_none; +} +MP_DEFINE_CONST_FUN_OBJ_2(displayio_group_remove_obj, displayio_group_obj_remove); + //| .. method:: __len__() //| //| Returns the number of layers in a Group @@ -281,7 +310,9 @@ STATIC const mp_rom_map_elem_t displayio_group_locals_dict_table[] = { { MP_ROM_QSTR(MP_QSTR_y), MP_ROM_PTR(&displayio_group_y_obj) }, { MP_ROM_QSTR(MP_QSTR_append), MP_ROM_PTR(&displayio_group_append_obj) }, { MP_ROM_QSTR(MP_QSTR_insert), MP_ROM_PTR(&displayio_group_insert_obj) }, + { MP_ROM_QSTR(MP_QSTR_index), MP_ROM_PTR(&displayio_group_index_obj) }, { MP_ROM_QSTR(MP_QSTR_pop), MP_ROM_PTR(&displayio_group_pop_obj) }, + { MP_ROM_QSTR(MP_QSTR_remove), MP_ROM_PTR(&displayio_group_remove_obj) }, }; STATIC MP_DEFINE_CONST_DICT(displayio_group_locals_dict, displayio_group_locals_dict_table); @@ -291,5 +322,6 @@ const mp_obj_type_t displayio_group_type = { .make_new = displayio_group_make_new, .subscr = group_subscr, .unary_op = group_unary_op, + .getiter = mp_obj_new_generic_iterator, .locals_dict = (mp_obj_dict_t*)&displayio_group_locals_dict, }; diff --git a/shared-bindings/displayio/Group.h b/shared-bindings/displayio/Group.h index 0570735cff..dc9616f2e9 100644 --- a/shared-bindings/displayio/Group.h +++ b/shared-bindings/displayio/Group.h @@ -44,6 +44,7 @@ void common_hal_displayio_group_append(displayio_group_t* self, mp_obj_t layer); void common_hal_displayio_group_insert(displayio_group_t* self, size_t index, mp_obj_t layer); size_t common_hal_displayio_group_get_len(displayio_group_t* self); mp_obj_t common_hal_displayio_group_pop(displayio_group_t* self, size_t index); +mp_int_t common_hal_displayio_group_index(displayio_group_t* self, mp_obj_t layer); mp_obj_t common_hal_displayio_group_get(displayio_group_t* self, size_t index); void common_hal_displayio_group_set(displayio_group_t* self, size_t index, mp_obj_t layer); diff --git a/shared-module/displayio/Group.c b/shared-module/displayio/Group.c index 097edf33bd..40f112bf17 100644 --- a/shared-module/displayio/Group.c +++ b/shared-module/displayio/Group.c @@ -95,6 +95,15 @@ mp_obj_t common_hal_displayio_group_pop(displayio_group_t* self, size_t index) { return item; } +mp_int_t common_hal_displayio_group_index(displayio_group_t* self, mp_obj_t layer) { + for (size_t i = 0; i < self->size; i++) { + if (self->children[i].original == layer) { + return i; + } + } + return -1; +} + size_t common_hal_displayio_group_get_len(displayio_group_t* self) { return self->size; }