get slicing

FoamyGuy · FoamyGuy · commit 0563487433fa · 2022-11-30T19:16:20.000-06:00
diff --git a/shared-bindings/_pixelmap/PixelMap.c b/shared-bindings/_pixelmap/PixelMap.c
@@ -135,7 +135,7 @@ MP_PROPERTY_GETTER(pixelmap_pixelmap_byteorder_obj,
     (mp_obj_t)&pixelmap_pixelmap_get_byteorder);
 
 //|
-//|     def fill(self, color: PixelType, /) -> None:
+//|     def fill(self, color: PixelType) -> None:
 //|         """Fill all the pixels in the map with the given color"""
 STATIC mp_obj_t pixelmap_pixelmap_fill(const mp_obj_t self_in, const mp_obj_t color) {
     pixelmap_pixelmap_obj_t *self = MP_OBJ_TO_PTR(self_in);
@@ -146,7 +146,7 @@ STATIC mp_obj_t pixelmap_pixelmap_fill(const mp_obj_t self_in, const mp_obj_t co
 MP_DEFINE_CONST_FUN_OBJ_2(pixelmap_pixelmap_fill_obj, pixelmap_pixelmap_fill);
 
 //|
-//|     def indices(self, index: int, /) -> Tuple[int]:
+//|     def indices(self, index: int) -> Tuple[int]:
 //|         """Return the PixelBuf indices for a PixelMap index"""
 STATIC mp_obj_t pixelmap_pixelmap_indices(const mp_obj_t self_in, const mp_obj_t index) {
     pixelmap_pixelmap_obj_t *self = MP_OBJ_TO_PTR(self_in);
@@ -157,10 +157,14 @@ STATIC mp_obj_t pixelmap_pixelmap_indices(const mp_obj_t self_in, const mp_obj_t
 MP_DEFINE_CONST_FUN_OBJ_2(pixelmap_pixelmap_indices_obj, pixelmap_pixelmap_indices);
 
 
+//|     @overload
+//|     def __getitem__(self, index: slice) -> PixelReturnSequence:
+//|         """Retrieve the value of the underlying pixels."""
+//|         ...
+//|     @overload
 //|     def __getitem__(self, index: int) -> PixelReturnType:
-//|         """Retrieve the value of one of the underlying pixels at 'index'.
-//|
-//|         Note that slices are not supported by PixelMap.__getitem__"""
+//|         """Retrieve the value of one of the underlying pixels at 'index'."""
+//|         ...
 //|     @overload
 //|     def __setitem__(self, index: slice, value: PixelSequence) -> None: ...
 //|     @overload
@@ -176,21 +180,41 @@ STATIC mp_obj_t pixelmap_pixelmap_subscr(mp_obj_t self_in, mp_obj_t index_in, mp
     if (value == MP_OBJ_NULL) {
         // delete
         return MP_OBJ_NULL; // op not supported
-    } else if (value == MP_OBJ_SENTINEL) {
-        int index = mp_obj_get_int(index_in);
-        return shared_module_pixelmap_pixelmap_getitem(self, index);
     }
 
-    // get
     if (0) {
     #if MICROPY_PY_BUILTINS_SLICE
     } else if (mp_obj_is_type(index_in, &mp_type_slice)) {
-        shared_module_pixelmap_pixelmap_setslice(self, index_in, value);
+        mp_bound_slice_t slice;
+        mp_seq_get_fast_slice_indexes(self->len, index_in, &slice);
+        size_t slice_len;
+        if (slice.step > 0) {
+            slice_len = slice.stop - slice.start;
+        } else {
+            slice_len = 1 + slice.start - slice.stop;
+        }
+        if (slice.step > 1 || slice.step < -1) {
+            size_t step = slice.step > 0 ? slice.step : slice.step * -1;
+            slice_len = (slice_len / step) + (slice_len % step ? 1 : 0);
+        }
+
+        if (value == MP_OBJ_SENTINEL) { // Get
+            return shared_module_pixelmap_pixelmap_getslice(self, slice, slice_len);
+        } else { // Set
+            shared_module_pixelmap_pixelmap_setslice(self, value, slice, slice_len);
+            return mp_const_none;
+        }
     #endif
-    } else {
-        shared_module_pixelmap_pixelmap_setitem(self, mp_obj_get_int(index_in), value);
+    } else { // single index
+        int index = mp_obj_get_int(index_in);
+
+        if (value == MP_OBJ_SENTINEL) { // Get
+            return shared_module_pixelmap_pixelmap_getitem(self, index);
+        } else {
+            shared_module_pixelmap_pixelmap_setitem(self, mp_obj_get_int(index_in), value);
+            return mp_const_none;
+        }
     }
-    return mp_const_none;
 }
 
 //|     def __len__(self) -> int:
diff --git a/shared-bindings/_pixelmap/PixelMap.h b/shared-bindings/_pixelmap/PixelMap.h
@@ -37,6 +37,7 @@ bool shared_module_pixelmap_pixelmap_auto_write_get(pixelmap_pixelmap_obj_t *sel
 void shared_module_pixelmap_pixelmap_auto_write_set(pixelmap_pixelmap_obj_t *self, bool auto_write);
 void shared_module_pixelmap_pixelmap_fill(pixelmap_pixelmap_obj_t *self, const mp_obj_t color);
 mp_obj_t shared_module_pixelmap_pixelmap_indices(pixelmap_pixelmap_obj_t *self, int index);
-void shared_module_pixelmap_pixelmap_setslice(pixelmap_pixelmap_obj_t *self, const mp_obj_t slice_in, const mp_obj_t value);
+void shared_module_pixelmap_pixelmap_setslice(pixelmap_pixelmap_obj_t *self, const mp_obj_t value, mp_bound_slice_t slice, size_t slice_len);
+mp_obj_t shared_module_pixelmap_pixelmap_getslice(pixelmap_pixelmap_obj_t *self, mp_bound_slice_t slice, size_t slice_len);
 mp_obj_t shared_module_pixelmap_pixelmap_getitem(pixelmap_pixelmap_obj_t *self, mp_int_t index);
 void shared_module_pixelmap_pixelmap_setitem(pixelmap_pixelmap_obj_t *self, mp_int_t index, const mp_obj_t value);
diff --git a/shared-module/_pixelmap/PixelMap.c b/shared-module/_pixelmap/PixelMap.c
@@ -34,7 +34,7 @@
 
 
 static void pixelmap_set_pixel_rgbw(pixelmap_pixelmap_obj_t *self, size_t i, color_u rgbw) {
-    mp_arg_validate_index_range(i, 0, self->len, MP_QSTR_index);
+    mp_arg_validate_index_range(i, 0, self->len - 1, MP_QSTR_index);
 
     mp_obj_t item = self->items[i];
     if (mp_obj_is_small_int(item)) {
@@ -109,20 +109,16 @@ mp_obj_t shared_module_pixelmap_pixelmap_indices(pixelmap_pixelmap_obj_t *self,
 }
 
 #if MICROPY_PY_BUILTINS_SLICE
-void shared_module_pixelmap_pixelmap_setslice(pixelmap_pixelmap_obj_t *self, const mp_obj_t slice_in, const mp_obj_t values) {
-    mp_bound_slice_t slice;
-    mp_seq_get_fast_slice_indexes(self->len, slice_in, &slice);
-    size_t slice_len;
-    if (slice.step > 0) {
-        slice_len = slice.stop - slice.start;
-    } else {
-        slice_len = 1 + slice.start - slice.stop;
-    }
-    if (slice.step > 1 || slice.step < -1) {
-        size_t step = slice.step > 0 ? slice.step : slice.step * -1;
-        slice_len = (slice_len / step) + (slice_len % step ? 1 : 0);
+
+mp_obj_t shared_module_pixelmap_pixelmap_getslice(pixelmap_pixelmap_obj_t *self, mp_bound_slice_t slice, size_t slice_len) {
+    mp_obj_tuple_t *t = MP_OBJ_TO_PTR(mp_obj_new_tuple(slice_len, NULL));
+    for (uint i = 0; i < slice_len; i++) {
+        t->items[i] = shared_module_pixelmap_pixelmap_getitem(self, i * slice.step + slice.start);
     }
+    return MP_OBJ_FROM_PTR(t);
+}
 
+void shared_module_pixelmap_pixelmap_setslice(pixelmap_pixelmap_obj_t *self, const mp_obj_t values, mp_bound_slice_t slice, size_t slice_len) {
     size_t num_items = mp_obj_get_int(mp_obj_len(values));
     if (num_items != slice_len) {
         mp_raise_ValueError_varg(translate("Unmatched number of items on RHS (expected %d, got %d)."), slice_len, num_items);
@@ -165,7 +161,7 @@ void shared_module_pixelmap_pixelmap_setitem(pixelmap_pixelmap_obj_t *self, mp_i
 }
 
 mp_obj_t shared_module_pixelmap_pixelmap_getitem(pixelmap_pixelmap_obj_t *self, mp_int_t i) {
-    mp_arg_validate_index_range(i, 0, self->len, MP_QSTR_index);
+    mp_arg_validate_index_range(i, 0, self->len - 1, MP_QSTR_index);
     mp_obj_t item = self->items[i];
     if (mp_obj_is_small_int(item)) {
         return common_hal_adafruit_pixelbuf_pixelbuf_get_pixel(self->pixelbuf, MP_OBJ_SMALL_INT_VALUE(item));
