Merge pull request #226 from pimoroni/patch-ws2812-rgbw

Plasma: Add WS2812 W suppot for #220
pimoroni · Dec 8, 2021 · 33ca1b6 · 33ca1b6
2 parents ce1f8ad + 580d8cf
commit 33ca1b6
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 14 deletions.
diff --git a/drivers/plasma/ws2812.cpp b/drivers/plasma/ws2812.cpp
@@ -67,7 +67,7 @@ void WS2812::clear() {
     }
 }
 
-void WS2812::set_hsv(uint32_t index, float h, float s, float v) {
+void WS2812::set_hsv(uint32_t index, float h, float s, float v, uint8_t w) {
     float i = floor(h * 6.0f);
     float f = h * 6.0f - i;
     v *= 255.0f;
@@ -76,12 +76,12 @@ void WS2812::set_hsv(uint32_t index, float h, float s, float v) {
     uint8_t t = v * (1.0f - (1.0f - f) * s);
 
     switch (int(i) % 6) {
-      case 0: set_rgb(index, v, t, p); break;
-      case 1: set_rgb(index, q, v, p); break;
-      case 2: set_rgb(index, p, v, t); break;
-      case 3: set_rgb(index, p, q, v); break;
-      case 4: set_rgb(index, t, p, v); break;
-      case 5: set_rgb(index, v, p, q); break;
+      case 0: set_rgb(index, v, t, p, w); break;
+      case 1: set_rgb(index, q, v, p, w); break;
+      case 2: set_rgb(index, p, v, t, w); break;
+      case 3: set_rgb(index, p, q, v, w); break;
+      case 4: set_rgb(index, t, p, v, w); break;
+      case 5: set_rgb(index, v, p, q, w); break;
     }
 }
 

diff --git a/drivers/plasma/ws2812.hpp b/drivers/plasma/ws2812.hpp
@@ -84,7 +84,7 @@ namespace plasma {
             bool stop();
             void update(bool blocking=false);
             void clear();
-            void set_hsv(uint32_t index, float h, float s, float v);
+            void set_hsv(uint32_t index, float h, float s, float v, uint8_t w=0);
             void set_rgb(uint32_t index, uint8_t r, uint8_t g, uint8_t b, uint8_t w=0, bool gamma=true);
             void set_brightness(uint8_t b);
             RGB get(uint32_t index) {return buffer[index];};

diff --git a/micropython/modules/plasma/plasma.cpp b/micropython/modules/plasma/plasma.cpp
@@ -133,13 +133,14 @@ mp_obj_t PlasmaWS2812_start(size_t n_args, const mp_obj_t *pos_args, mp_map_t *k
 
 
 mp_obj_t PlasmaWS2812_set_rgb(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
-    enum { ARG_self, ARG_index, ARG_r, ARG_g, ARG_b };
+    enum { ARG_self, ARG_index, ARG_r, ARG_g, ARG_b, ARG_w };
     static const mp_arg_t allowed_args[] = {
         { MP_QSTR_, MP_ARG_REQUIRED | MP_ARG_OBJ },
         { MP_QSTR_index, MP_ARG_REQUIRED | MP_ARG_INT },
         { MP_QSTR_r, MP_ARG_REQUIRED | MP_ARG_INT },
         { MP_QSTR_g, MP_ARG_REQUIRED | MP_ARG_INT },
-        { MP_QSTR_b, MP_ARG_REQUIRED | MP_ARG_INT }
+        { MP_QSTR_b, MP_ARG_REQUIRED | MP_ARG_INT },
+        { MP_QSTR_w, MP_ARG_INT, {.u_int = 0} }
     };
 
     // Parse args.
@@ -150,21 +151,24 @@ mp_obj_t PlasmaWS2812_set_rgb(size_t n_args, const mp_obj_t *pos_args, mp_map_t
     int r = args[ARG_r].u_int;
     int g = args[ARG_g].u_int;
     int b = args[ARG_b].u_int;
+    int w = args[ARG_w].u_int;
 
     _PlasmaWS2812_obj_t *self = MP_OBJ_TO_PTR2(args[ARG_self].u_obj, _PlasmaWS2812_obj_t);
-    self->ws2812->set_rgb(index, r, g, b);
+
+    self->ws2812->set_rgb(index, r, g, b, w);
 
     return mp_const_none;
 }
 
 mp_obj_t PlasmaWS2812_set_hsv(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
-    enum { ARG_self, ARG_index, ARG_h, ARG_s, ARG_v };
+    enum { ARG_self, ARG_index, ARG_h, ARG_s, ARG_v, ARG_w };
     static const mp_arg_t allowed_args[] = {
         { MP_QSTR_, MP_ARG_REQUIRED | MP_ARG_OBJ },
         { MP_QSTR_index, MP_ARG_REQUIRED | MP_ARG_INT },
         { MP_QSTR_hue, MP_ARG_REQUIRED | MP_ARG_OBJ },
         { MP_QSTR_sat, MP_ARG_OBJ, {.u_rom_obj = MP_ROM_PTR(&const_float_1)} },
-        { MP_QSTR_val, MP_ARG_OBJ, {.u_rom_obj = MP_ROM_PTR(&const_float_1)} }
+        { MP_QSTR_val, MP_ARG_OBJ, {.u_rom_obj = MP_ROM_PTR(&const_float_1)} },
+        { MP_QSTR_w, MP_ARG_INT, {.u_int = 0} }
     };
 
     // Parse args.
@@ -175,9 +179,11 @@ mp_obj_t PlasmaWS2812_set_hsv(size_t n_args, const mp_obj_t *pos_args, mp_map_t
     float h = mp_obj_get_float(args[ARG_h].u_obj);
     float s = mp_obj_get_float(args[ARG_s].u_obj);
     float v = mp_obj_get_float(args[ARG_v].u_obj);
+    int w = args[ARG_w].u_int;
 
     _PlasmaWS2812_obj_t *self = MP_OBJ_TO_PTR2(args[ARG_self].u_obj, _PlasmaWS2812_obj_t);
-    self->ws2812->set_hsv(index, h, s, v);
+
+    self->ws2812->set_hsv(index, h, s, v, w);
 
     return mp_const_none;
 }