PixelBuf: Add support for 2D buffers

Author: jepler

shared-bindings/_pixelbuf/PixelBuf.c

@@ -72,14 +72,15 @@ static void parse_byteorder(mp_obj_t byteorder_obj, pixelbuf_byteorder_details_t
 //|
 STATIC mp_obj_t pixelbuf_pixelbuf_make_new(const mp_obj_type_t *type, size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
     mp_arg_check_num(n_args, kw_args, 1, MP_OBJ_FUN_ARGS_MAX, true);
-    enum { ARG_size, ARG_byteorder, ARG_brightness, ARG_auto_write, ARG_header, ARG_trailer };
+    enum { ARG_size, ARG_byteorder, ARG_brightness, ARG_auto_write, ARG_header, ARG_trailer, ARG_width };
     static const mp_arg_t allowed_args[] = {
         { MP_QSTR_size, MP_ARG_REQUIRED | MP_ARG_INT },
         { MP_QSTR_byteorder, MP_ARG_KW_ONLY | MP_ARG_OBJ, { .u_obj = MP_OBJ_NEW_QSTR(MP_QSTR_BGR) } },
         { MP_QSTR_brightness, MP_ARG_KW_ONLY | MP_ARG_OBJ, { .u_obj = mp_const_none } },
         { MP_QSTR_auto_write, MP_ARG_KW_ONLY | MP_ARG_BOOL, {.u_bool = false} },
         { MP_QSTR_header, MP_ARG_KW_ONLY | MP_ARG_OBJ, { .u_obj = mp_const_none } },
         { MP_QSTR_trailer, MP_ARG_KW_ONLY | MP_ARG_OBJ, { .u_obj = mp_const_none } },
+        { MP_QSTR_width, MP_ARG_KW_ONLY | MP_ARG_OBJ, { .u_obj = mp_const_none } },
     };
     mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
     mp_arg_parse_all(n_args, pos_args, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
@@ -89,7 +90,11 @@ STATIC mp_obj_t pixelbuf_pixelbuf_make_new(const mp_obj_type_t *type, size_t n_a
 
     mp_buffer_info_t header_bufinfo;
     mp_buffer_info_t trailer_bufinfo;
+    int width = -1;
 
+    if (args[ARG_width].u_obj != mp_const_none) {
+        width = MP_OBJ_SMALL_INT_VALUE(args[ARG_width].u_obj);
+    }
     if (!mp_get_buffer(args[ARG_header].u_obj, &header_bufinfo, MP_BUFFER_READ)) {
         header_bufinfo.buf = NULL;
         header_bufinfo.len = 0;
@@ -114,7 +119,7 @@ STATIC mp_obj_t pixelbuf_pixelbuf_make_new(const mp_obj_type_t *type, size_t n_a
     self->base.type = &pixelbuf_pixelbuf_type;
     common_hal__pixelbuf_pixelbuf_construct(self, args[ARG_size].u_int,
     &byteorder_details, brightness, args[ARG_auto_write].u_bool, header_bufinfo.buf,
-    header_bufinfo.len, trailer_bufinfo.buf, trailer_bufinfo.len);
+    header_bufinfo.len, trailer_bufinfo.buf, trailer_bufinfo.len, width);
 
     return MP_OBJ_FROM_PTR(self);
 }
@@ -277,6 +282,27 @@ STATIC mp_obj_t pixelbuf_pixelbuf_fill(mp_obj_t self_in, mp_obj_t value) {
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(pixelbuf_pixelbuf_fill_obj, pixelbuf_pixelbuf_fill);
 
+STATIC size_t pixelbuf_get_index(mp_obj_t self_in, mp_obj_t index_in) {
+    size_t length = common_hal__pixelbuf_pixelbuf_get_len(self_in);
+    int16_t width = common_hal__pixelbuf_pixelbuf_get_width(self_in);
+    if(width != -1 && MP_OBJ_IS_TYPE(index_in, &mp_type_tuple)) {
+        size_t x = mp_get_index(mp_obj_get_type(self_in),
+                                    width,
+                                    mp_obj_tuple_subscr(index_in,
+                                                        MP_OBJ_NEW_SMALL_INT(0),
+                                                        MP_OBJ_SENTINEL),
+                                    false);
+        size_t y = mp_get_index(mp_obj_get_type(self_in),
+                                    common_hal__pixelbuf_pixelbuf_get_height(self_in),
+                                    mp_obj_tuple_subscr(index_in,
+                                                        MP_OBJ_NEW_SMALL_INT(1),
+                                                        MP_OBJ_SENTINEL),
+                                    false);
+        return x + width * y;
+    } else {
+        return mp_get_index(self_in, length, index_in, false);
+    }
+}
 
 //|   .. method:: __getitem__(index)
 //|
@@ -353,8 +379,7 @@ STATIC mp_obj_t pixelbuf_pixelbuf_subscr(mp_obj_t self_in, mp_obj_t index_in, mp
         }
 #endif
     } else { // Single index rather than slice.
-        size_t length = common_hal__pixelbuf_pixelbuf_get_len(self_in);
-        size_t index = mp_get_index(mp_obj_get_type(self_in), length, index_in, false);
+        size_t index = pixelbuf_get_index(self_in, index_in);
 
         if (value == MP_OBJ_SENTINEL) { // Get
             return common_hal__pixelbuf_pixelbuf_get_pixel(self_in, index);

shared-bindings/_pixelbuf/PixelBuf.h

@@ -33,7 +33,7 @@ const mp_obj_type_t pixelbuf_pixelbuf_type;
 
 void common_hal__pixelbuf_pixelbuf_construct(pixelbuf_pixelbuf_obj_t *self, size_t n,
     pixelbuf_byteorder_details_t* byteorder, mp_float_t brightness, bool auto_write, uint8_t* header,
-    size_t header_len, uint8_t* trailer, size_t trailer_len);
+    size_t header_len, uint8_t* trailer, size_t trailer_len, int width);
 
 // These take mp_obj_t because they are called on subclasses of PixelBuf.
 uint8_t common_hal__pixelbuf_pixelbuf_get_bpp(mp_obj_t self);
@@ -42,6 +42,8 @@ void common_hal__pixelbuf_pixelbuf_set_brightness(mp_obj_t self, mp_float_t brig
 bool common_hal__pixelbuf_pixelbuf_get_auto_write(mp_obj_t self);
 void common_hal__pixelbuf_pixelbuf_set_auto_write(mp_obj_t self, bool auto_write);
 size_t common_hal__pixelbuf_pixelbuf_get_len(mp_obj_t self_in);
+int common_hal__pixelbuf_pixelbuf_get_width(mp_obj_t self_in);
+int common_hal__pixelbuf_pixelbuf_get_height(mp_obj_t self_in);
 mp_obj_t common_hal__pixelbuf_pixelbuf_get_byteorder_string(mp_obj_t self);
 void common_hal__pixelbuf_pixelbuf_fill(mp_obj_t self, mp_obj_t item);
 void common_hal__pixelbuf_pixelbuf_show(mp_obj_t self);

shared-module/_pixelbuf/PixelBuf.c

@@ -41,12 +41,13 @@ static pixelbuf_pixelbuf_obj_t* native_pixelbuf(mp_obj_t pixelbuf_obj) {
 
 void common_hal__pixelbuf_pixelbuf_construct(pixelbuf_pixelbuf_obj_t *self, size_t n,
         pixelbuf_byteorder_details_t* byteorder, mp_float_t brightness, bool auto_write,
-        uint8_t* header, size_t header_len, uint8_t* trailer, size_t trailer_len) {
+        uint8_t* header, size_t header_len, uint8_t* trailer, size_t trailer_len, int width) {
 
     self->pixel_count = n;
     self->byteorder = *byteorder;  // Copied because we modify for dotstar
     self->bytes_per_pixel = byteorder->is_dotstar ? 4 : byteorder->bpp;
     self->auto_write = false;
+    self->width = width;
 
     size_t pixel_len = self->pixel_count * self->bytes_per_pixel;
     self->transmit_buffer_obj = mp_obj_new_bytes_of_zeros(header_len + pixel_len + trailer_len);
@@ -80,6 +81,16 @@ size_t common_hal__pixelbuf_pixelbuf_get_len(mp_obj_t self_in) {
     return self->pixel_count;
 }
 
+int common_hal__pixelbuf_pixelbuf_get_width(mp_obj_t self_in) {
+    pixelbuf_pixelbuf_obj_t* self = native_pixelbuf(self_in);
+    return self->width;
+}
+
+int common_hal__pixelbuf_pixelbuf_get_height(mp_obj_t self_in) {
+    pixelbuf_pixelbuf_obj_t* self = native_pixelbuf(self_in);
+    return self->width == -1 ? -1 : (int)self->pixel_count / self->width;
+}
+
 uint8_t common_hal__pixelbuf_pixelbuf_get_bpp(mp_obj_t self_in) {
     pixelbuf_pixelbuf_obj_t* self = native_pixelbuf(self_in);
     return self->byteorder.bpp;

shared-module/_pixelbuf/PixelBuf.h

@@ -50,6 +50,7 @@ typedef struct {
     mp_obj_base_t base;
     size_t pixel_count;
     size_t bytes_per_pixel;
+    int16_t width;
     pixelbuf_byteorder_details_t byteorder;
     mp_float_t brightness;
     mp_obj_t transmit_buffer_obj;