neopixel: Add support for RGBW LEDs

Signed-off-by: Kevin O'Connor <kevin@koconnor.net>

4 files changed, 36 insertions, 17 deletions

diff --git a/config/example-extras.cfg b/config/example-extras.cfg
index 2291a915..854d38ff 100644
--- a/config/example-extras.cfg
+++ b/config/example-extras.cfg
@@ -1354,14 +1354,16 @@
 #   The number of Neopixel chips that are "daisy chained" to the
 #   provided pin. The default is 1 (which indicates only a single
 #   Neopixel is connected to the pin).
-#color_order_GRB: True
-#   Set the pixel order to green, red, blue. If using the WS2811 chip
-#   (in 800Khz mode) then set this to False. The default is True.
+#color_order: GRB
+#   Set the pixel order required by the LED hardware. Options are GRB,
+#   RGB, GRBW, or RGBW. The default is GRB.
 #initial_RED: 0.0
 #initial_GREEN: 0.0
 #initial_BLUE: 0.0
+#initial_WHITE: 0.0
 #   Sets the initial LED color of the Neopixel. Each value should be
-#   between 0.0 and 1.0. The default for each color is 0.
+#   between 0.0 and 1.0. The WHITE option is only available on RGBW
+#   LEDs. The default for each color is 0.
 
 # Dotstar (aka APA102) LED support (one may define any number of
 # sections with a "dotstar" prefix). One may set the LED color via
diff --git a/docs/Config_Changes.md b/docs/Config_Changes.md
index ba13ffc0..9153f64a 100644
--- a/docs/Config_Changes.md
+++ b/docs/Config_Changes.md
@@ -6,6 +6,10 @@ All dates in this document are approximate.
 
 # Changes
 
+20201029: The neopixel `color_order_GRB` config option has been
+removed. If necessary, update the config to set the new `color_order`
+option to RGB, GRB, RGBW, or GRBW.
+
 20201029: The serial option in the mcu config section no longer
 defaults to /dev/ttyS0.  In the rare situation where /dev/ttyS0 is the
 desired serial port, it must be specified explicitly.
diff --git a/docs/G-Codes.md b/docs/G-Codes.md
index b2a35b52..43a3ae85 100644
--- a/docs/G-Codes.md
+++ b/docs/G-Codes.md
@@ -259,8 +259,9 @@ is enabled:
 The following command is available when "neopixel" or "dotstar" config
 sections are enabled:
 - `SET_LED LED=<config_name> RED=<value> GREEN=<value> BLUE=<value>
-  [INDEX=<index>] [TRANSMIT=0]`: This sets the LED output. Each color
-  <value> must be between 0.0 and 1.0. If multiple LED chips are
+  WHITE=<value> [INDEX=<index>] [TRANSMIT=0]`: This sets the LED
+  output. Each color <value> must be between 0.0 and 1.0. The WHITE
+  option is only valid on RGBW LEDs. If multiple LED chips are
   daisy-chained then one may specify INDEX to alter the color of just
   the given chip (1 for the first chip, 2 for the second, etc.). If
   INDEX is not provided then all LEDs in the daisy-chain will be set
diff --git a/klippy/extras/neopixel.py b/klippy/extras/neopixel.py
index 307651ea..422b849d 100644
--- a/klippy/extras/neopixel.py
+++ b/klippy/extras/neopixel.py
@@ -23,16 +23,21 @@ class PrinterNeoPixel:
         self.oid = self.mcu.create_oid()
         self.pin = pin_params['pin']
         self.mcu.register_config_callback(self.build_config)
-        self.color_order_GRB = config.getboolean("color_order_GRB", True)
-        self.chain_count = config.getint('chain_count', 1,
-                                         minval=1, maxval=MAX_MCU_SIZE//3)
+        formats = {v: v for v in ["RGB", "GRB", "RGBW", "GRBW"]}
+        self.color_order = config.getchoice("color_order", formats, "GRB")
+        elem_size = len(self.color_order)
+        self.chain_count = config.getint('chain_count', 1, minval=1,
+                                         maxval=MAX_MCU_SIZE//elem_size)
         self.neopixel_update_cmd = self.neopixel_send_cmd = None
         # Initial color
-        self.color_data = bytearray(self.chain_count * 3)
+        self.color_data = bytearray(self.chain_count * elem_size)
         red = config.getfloat('initial_RED', 0., minval=0., maxval=1.)
         green = config.getfloat('initial_GREEN', 0., minval=0., maxval=1.)
         blue = config.getfloat('initial_BLUE', 0., minval=0., maxval=1.)
-        self.update_color_data(red, green, blue)
+        white = 0
+        if elem_size == 4:
+            white = config.getfloat('initial_WHITE', 0., minval=0., maxval=1.)
+        self.update_color_data(red, green, blue, white)
         self.old_color_data = bytearray([d ^ 1 for d in self.color_data])
         # Register commands
         self.printer.register_event_handler("klippy:connect", self.send_data)
@@ -44,25 +49,31 @@ class PrinterNeoPixel:
         rmt = self.mcu.seconds_to_clock(RESET_MIN_TIME)
         self.mcu.add_config_cmd("config_neopixel oid=%d pin=%s data_size=%d"
                                 " bit_max_ticks=%d reset_min_ticks=%d"
-                                % (self.oid, self.pin, self.chain_count * 3,
+                                % (self.oid, self.pin, len(self.color_data),
                                    bmt, rmt))
         cmd_queue = self.mcu.alloc_command_queue()
         self.neopixel_update_cmd = self.mcu.lookup_command(
             "neopixel_update oid=%c pos=%hu data=%*s", cq=cmd_queue)
         self.neopixel_send_cmd = self.mcu.lookup_query_command(
             "neopixel_send oid=%c", "neopixel_result success=%c", cq=cmd_queue)
-    def update_color_data(self, red, green, blue, index=None):
+    def update_color_data(self, red, green, blue, white, index=None):
         red = int(red * 255. + .5)
         blue = int(blue * 255. + .5)
         green = int(green * 255. + .5)
-        if self.color_order_GRB:
+        white = int(white * 255. + .5)
+        if self.color_order == "GRB":
             color_data = [green, red, blue]
-        else:
+        elif self.color_order == "RGB":
             color_data = [red, green, blue]
+        elif self.color_order == "GRBW":
+            color_data = [green, red, blue, white]
+        else:
+            color_data = [red, green, blue, white]
         if index is None:
             self.color_data[:] = color_data * self.chain_count
         else:
-            self.color_data[(index-1)*3:index*3] = color_data
+            elem_size = len(color_data)
+            self.color_data[(index-1)*elem_size:index*elem_size] = color_data
     def send_data(self, minclock=0):
         old_data, new_data = self.old_color_data, self.color_data
         if new_data == old_data:
@@ -97,9 +108,10 @@ class PrinterNeoPixel:
         red = gcmd.get_float('RED', 0., minval=0., maxval=1.)
         green = gcmd.get_float('GREEN', 0., minval=0., maxval=1.)
         blue = gcmd.get_float('BLUE', 0., minval=0., maxval=1.)
+        white = gcmd.get_float('WHITE', 0., minval=0., maxval=1.)
         index = gcmd.get_int('INDEX', None, minval=1, maxval=self.chain_count)
         transmit = gcmd.get_int('TRANSMIT', 1)
-        self.update_color_data(red, green, blue, index)
+        self.update_color_data(red, green, blue, white, index)
         # Send command
         if not transmit:
             return