make pixelbuf handle floats and iterables

2 files changed, 12 insertions, 20 deletions

diff --git a/shared-bindings/_pixelbuf/PixelBuf.c b/shared-bindings/_pixelbuf/PixelBuf.c
index 4fc6f7cbf..ee7dac4b3 100644
--- a/shared-bindings/_pixelbuf/PixelBuf.c
+++ b/shared-bindings/_pixelbuf/PixelBuf.c
@@ -40,6 +40,10 @@
 #include "shared-module/_pixelbuf/PixelBuf.h"
 #include "shared-bindings/digitalio/DigitalInOut.h"
 
+#ifdef CIRCUITPY_ULAB
+#include "extmod/ulab/code/ndarray.h"
+#endif
+
 extern const int32_t colorwheel(float pos);
 
 static void parse_byteorder(mp_obj_t byteorder_obj, pixelbuf_byteorder_details_t* parsed);
@@ -326,27 +330,15 @@ STATIC mp_obj_t pixelbuf_pixelbuf_subscr(mp_obj_t self_in, mp_obj_t index_in, mp
         } else { // Set
             #if MICROPY_PY_ARRAY_SLICE_ASSIGN
 
-            if (!(MP_OBJ_IS_TYPE(value, &mp_type_list) || MP_OBJ_IS_TYPE(value, &mp_type_tuple))) {
-                mp_raise_ValueError(translate("tuple/list required on RHS"));
-            }
-
-            mp_obj_t *src_objs;
-            size_t num_items;
-            if (MP_OBJ_IS_TYPE(value, &mp_type_list)) {
-                mp_obj_list_t *t = MP_OBJ_TO_PTR(value);
-                num_items = t->len;
-                src_objs = t->items;
-            } else {
-                mp_obj_tuple_t *l = MP_OBJ_TO_PTR(value);
-                num_items = l->len;
-                src_objs = l->items;
-            }
-            if (num_items != slice_len) {
-                mp_raise_ValueError_varg(translate("Unmatched number of items on RHS (expected %d, got %d)."),
-                                                   slice_len, num_items);
+            size_t num_items = mp_obj_get_int(mp_obj_len(value));
+            
+            if (num_items != slice_len && num_items != (slice_len * common_hal__pixelbuf_pixelbuf_get_bpp(self_in))) {
+                mp_raise_ValueError_varg(translate("Unmatched number of items on RHS (expected %d or %d, got %d)."),
+                                                   slice_len, slice_len * common_hal__pixelbuf_pixelbuf_get_bpp(self_in),
+                                                   num_items);
             }
 
-            common_hal__pixelbuf_pixelbuf_set_pixels(self_in, slice.start, slice.step, slice_len, src_objs);
+            common_hal__pixelbuf_pixelbuf_set_pixels(self_in, slice.start, slice.step, slice_len, value, num_items != slice_len);
             return mp_const_none;
             #else
             return MP_OBJ_NULL; // op not supported
diff --git a/shared-bindings/_pixelbuf/PixelBuf.h b/shared-bindings/_pixelbuf/PixelBuf.h
index 0b09a5771..111223be0 100644
--- a/shared-bindings/_pixelbuf/PixelBuf.h
+++ b/shared-bindings/_pixelbuf/PixelBuf.h
@@ -47,6 +47,6 @@ void common_hal__pixelbuf_pixelbuf_fill(mp_obj_t self, mp_obj_t item);
 void common_hal__pixelbuf_pixelbuf_show(mp_obj_t self);
 mp_obj_t common_hal__pixelbuf_pixelbuf_get_pixel(mp_obj_t self, size_t index);
 void common_hal__pixelbuf_pixelbuf_set_pixel(mp_obj_t self, size_t index, mp_obj_t item);
-void common_hal__pixelbuf_pixelbuf_set_pixels(mp_obj_t self_in, size_t start, mp_int_t step, size_t slice_len, mp_obj_t* values);
+void common_hal__pixelbuf_pixelbuf_set_pixels(mp_obj_t self_in, size_t start, mp_int_t step, size_t slice_len, mp_obj_t* values, bool flattened);
 
 #endif  // CP_SHARED_BINDINGS_PIXELBUF_PIXELBUF_H