Merge pull request #3260 from rhooper/pixelbuf-iterable

Make pixelbuf support iterables and floats

5 files changed, 61 insertions, 34 deletions

diff --git a/locale/circuitpython.pot b/locale/circuitpython.pot
index a1b286e2e..ffb2e91ef 100644
--- a/locale/circuitpython.pot
+++ b/locale/circuitpython.pot
@@ -1958,6 +1958,7 @@ msgid "can't assign to expression"
 msgstr ""
 
 #: py/obj.c py/objint.c shared-bindings/i2cperipheral/I2CPeripheral.c
+#: shared-module/_pixelbuf/PixelBuf.c
 msgid "can't convert %q to %q"
 msgstr ""
 
@@ -3155,10 +3156,6 @@ msgstr ""
 msgid "tuple/list has wrong length"
 msgstr ""
 
-#: shared-bindings/_pixelbuf/PixelBuf.c
-msgid "tuple/list required on RHS"
-msgstr ""
-
 #: ports/atmel-samd/common-hal/busio/UART.c ports/nrf/common-hal/busio/UART.c
 #: shared-bindings/busio/UART.c
 msgid "tx and rx cannot both be None"
diff --git a/ports/atmel-samd/boards/bdmicro_vina_m0/mpconfigboard.mk b/ports/atmel-samd/boards/bdmicro_vina_m0/mpconfigboard.mk
index a36966b87..b7ad47ff2 100644
--- a/ports/atmel-samd/boards/bdmicro_vina_m0/mpconfigboard.mk
+++ b/ports/atmel-samd/boards/bdmicro_vina_m0/mpconfigboard.mk
@@ -15,5 +15,5 @@ CIRCUITPY_BITBANGIO = 0
 CIRCUITPY_I2CPERIPHERAL = 0
 CIRCUITPY_VECTORIO = 0
 
-CFLAGS_INLINE_LIMIT = 60
+CFLAGS_INLINE_LIMIT = 50
 SUPEROPT_GC = 0
diff --git a/shared-bindings/_pixelbuf/PixelBuf.c b/shared-bindings/_pixelbuf/PixelBuf.c
index 4fc6f7cbf..88e8aa801 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);
@@ -305,6 +309,16 @@ STATIC mp_obj_t pixelbuf_pixelbuf_subscr(mp_obj_t self_in, mp_obj_t index_in, mp
 
         size_t length = common_hal__pixelbuf_pixelbuf_get_len(self_in);
         mp_seq_get_fast_slice_indexes(length, index_in, &slice);
+        static mp_obj_tuple_t flat_item_tuple = {
+            .base = {&mp_type_tuple},
+            .len = 0,
+            .items = {
+                mp_const_none,
+                mp_const_none,
+                mp_const_none,
+                mp_const_none,
+            }
+        };
 
         size_t slice_len;
         if (slice.step > 0) {
@@ -326,27 +340,13 @@ 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"));
-            }
+            size_t num_items = mp_obj_get_int(mp_obj_len(value));
 
-            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);
+            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, got %d)."), slice_len, 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 ? &flat_item_tuple : mp_const_none);
             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..d41082059 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, mp_obj_tuple_t *flatten_to);
 
 #endif  // CP_SHARED_BINDINGS_PIXELBUF_PIXELBUF_H
diff --git a/shared-module/_pixelbuf/PixelBuf.c b/shared-module/_pixelbuf/PixelBuf.c
index 9d671e454..f3e679631 100644
--- a/shared-module/_pixelbuf/PixelBuf.c
+++ b/shared-module/_pixelbuf/PixelBuf.c
@@ -132,6 +132,18 @@ void common_hal__pixelbuf_pixelbuf_set_brightness(mp_obj_t self_in, mp_float_t b
     }
 }
 
+uint8_t _pixelbuf_get_as_uint8(mp_obj_t obj) {
+    if (MP_OBJ_IS_SMALL_INT(obj)) {
+        return MP_OBJ_SMALL_INT_VALUE(obj);
+    } else if (MP_OBJ_IS_INT(obj)) {
+        return mp_obj_get_int_truncated(obj);
+    } else if (mp_obj_is_float(obj)) {
+        return (uint8_t)mp_obj_get_float(obj);
+    }
+    mp_raise_TypeError_varg(
+            translate("can't convert %q to %q"), mp_obj_get_type_qstr(obj), MP_QSTR_int);
+}
+
 void _pixelbuf_parse_color(pixelbuf_pixelbuf_obj_t* self, mp_obj_t color, uint8_t* r, uint8_t* g, uint8_t* b, uint8_t* w) {
     pixelbuf_byteorder_details_t *byteorder = &self->byteorder;
     // w is shared between white in NeoPixels and brightness in dotstars (so that DotStars can have
@@ -142,8 +154,8 @@ void _pixelbuf_parse_color(pixelbuf_pixelbuf_obj_t* self, mp_obj_t color, uint8_
         *w = 0;
     }
 
-    if (MP_OBJ_IS_INT(color)) {
-        mp_int_t value = mp_obj_get_int_truncated(color);
+    if (MP_OBJ_IS_INT(color) || mp_obj_is_float(color)) {
+        mp_int_t value = MP_OBJ_IS_INT(color) ? mp_obj_get_int_truncated(color) : mp_obj_get_float(color);
         *r = value >> 16 & 0xff;
         *g = (value >> 8) & 0xff;
         *b = value & 0xff;
@@ -155,9 +167,9 @@ void _pixelbuf_parse_color(pixelbuf_pixelbuf_obj_t* self, mp_obj_t color, uint8_
             mp_raise_ValueError_varg(translate("Expected tuple of length %d, got %d"), byteorder->bpp, len);
         }
 
-        *r = mp_obj_get_int_truncated(items[PIXEL_R]);
-        *g = mp_obj_get_int_truncated(items[PIXEL_G]);
-        *b = mp_obj_get_int_truncated(items[PIXEL_B]);
+        *r = _pixelbuf_get_as_uint8(items[PIXEL_R]);
+        *g = _pixelbuf_get_as_uint8(items[PIXEL_G]);
+        *b = _pixelbuf_get_as_uint8(items[PIXEL_B]);
         if (len > 3) {
             if (mp_obj_is_float(items[PIXEL_W])) {
                 *w = 255 * mp_obj_get_float(items[PIXEL_W]);
@@ -218,17 +230,35 @@ void _pixelbuf_set_pixel(pixelbuf_pixelbuf_obj_t* self, size_t index, mp_obj_t v
     _pixelbuf_set_pixel_color(self, index, r, g, b, w);
 }
 
-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,
+    mp_obj_tuple_t *flatten_to)
+{
     pixelbuf_pixelbuf_obj_t* self = native_pixelbuf(self_in);
-    for (size_t i = 0; i < slice_len; i++) {
-        _pixelbuf_set_pixel(self, start, values[i]);
-        start+=step;
+    mp_obj_iter_buf_t iter_buf;
+    mp_obj_t iterable = mp_getiter(values, &iter_buf);
+    mp_obj_t item;
+    size_t i = 0;
+    bool flattened = flatten_to != mp_const_none;
+    if (flattened) flatten_to->len = self->bytes_per_pixel;
+    while ((item = mp_iternext(iterable)) != MP_OBJ_STOP_ITERATION) {
+        if (flattened) {
+            flatten_to->items[i % self->bytes_per_pixel] = item;
+            if (++i % self->bytes_per_pixel == 0) {
+                _pixelbuf_set_pixel(self, start, flatten_to);
+                start+=step;
+            }
+        } else {
+            _pixelbuf_set_pixel(self, start, item);
+            start+=step;
+        }
     }
     if (self->auto_write) {
         common_hal__pixelbuf_pixelbuf_show(self_in);
     }
 }
 
+
+
 void common_hal__pixelbuf_pixelbuf_set_pixel(mp_obj_t self_in, size_t index, mp_obj_t value) {
     pixelbuf_pixelbuf_obj_t* self = native_pixelbuf(self_in);
     _pixelbuf_set_pixel(self, index, value);