Merge pull request #1888 from tannewt/iterable_group4.1.0-alpha.0

Improvements to Group

5 files changed, 76 insertions, 0 deletions

diff --git a/py/obj.c b/py/obj.c
index fb59eec82..322a302f9 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 ae2bff190..773043b59 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 bbf7190a3..76719c580 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 0570735cf..dc9616f2e 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 097edf33b..40f112bf1 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;
 }