led: Fix SET_LED TRANSMIT=1 handling

If the final update in a series of SET_LED update commands didn't
alter a color than the transmit would not occur - even if prior
commands did make color changes.  Refactor the update code and fix.

Reported by @mstansberry.

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

1 files changed, 29 insertions, 25 deletions

diff --git a/klippy/extras/led.py b/klippy/extras/led.py
index 5830d77e..e1c0a822 100644
--- a/klippy/extras/led.py
+++ b/klippy/extras/led.py
@@ -15,6 +15,7 @@ class LEDHelper:
         self.printer = config.get_printer()
         self.update_func = update_func
         self.led_count = led_count
+        self.need_transmit = False
         # Initial color
         red = config.getfloat('initial_RED', 0., minval=0., maxval=1.)
         green = config.getfloat('initial_GREEN', 0., minval=0., maxval=1.)
@@ -28,6 +29,28 @@ class LEDHelper:
         gcode = self.printer.lookup_object('gcode')
         gcode.register_mux_command("SET_LED", "LED", name, self.cmd_SET_LED,
                                    desc=self.cmd_SET_LED_help)
+    def get_led_count(self):
+        return self.led_count
+    def set_color(self, index, color):
+        if index is None:
+            new_led_state = [color] * self.led_count
+            if self.led_state == new_led_state:
+                return
+        else:
+            if self.led_state[index - 1] == color:
+                return
+            new_led_state = list(self.led_state)
+            new_led_state[index - 1] = color
+        self.led_state = new_led_state
+        self.need_transmit = True
+    def check_transmit(self, print_time):
+        if not self.need_transmit:
+            return
+        self.need_transmit = False
+        try:
+            self.update_func(self.led_state, print_time)
+        except self.printer.command_error as e:
+            logging.exception("led update transmit error")
     cmd_SET_LED_help = "Set the color of an LED"
     def cmd_SET_LED(self, gcmd):
         # Parse parameters
@@ -41,21 +64,9 @@ class LEDHelper:
         color = (red, green, blue, white)
         # Update and transmit data
         def lookahead_bgfunc(print_time):
-            if index is None:
-                new_led_state = [color] * self.led_count
-                if self.led_state == new_led_state:
-                    return
-                self.led_state = new_led_state
-            else:
-                if self.led_state[index - 1] == color:
-                    return
-                self.led_state = led_state = list(self.led_state)
-                led_state[index - 1] = color
+            self.set_color(index, color)
             if transmit:
-                try:
-                    self.update_func(self.led_state, print_time)
-                except self.printer.command_error as e:
-                    logging.exception("led update transmit error")
+                self.check_transmit(print_time)
         if sync:
             #Sync LED Update with print time and send
             toolhead = self.printer.lookup_object('toolhead')
@@ -129,19 +140,12 @@ class PrinterLED:
                 if len(parts) < 4:
                     parts += [0.] * (4 - len(parts))
                 rendered[uid] = color = tuple(parts)
-            prev_color = led_helper.led_state[index-1]
-            if color != prev_color:
-                if led_helper not in need_transmit:
-                    need_transmit[led_helper] = 1
-                    led_helper.led_state = list(led_helper.led_state)
-                led_helper.led_state[index-1] = color
+            need_transmit[led_helper] = 1
+            led_helper.set_color(index, color)
         context.clear() # Remove circular references for better gc
         # Transmit pending changes
         for led_helper in need_transmit.keys():
-            try:
-                led_helper.update_func(led_helper.led_state, None)
-            except Exception as e:
-                logging.exception("led template transmit error")
+            led_helper.check_transmit(None)
         return eventtime + RENDER_TIME
     cmd_SET_LED_TEMPLATE_help = "Assign a display_template to an LED"
     def cmd_SET_LED_TEMPLATE(self, gcmd):
@@ -149,7 +153,7 @@ class PrinterLED:
         led_helper = self.led_helpers.get(led_name)
         if led_helper is None:
             raise gcmd.error("Unknown LED '%s'" % (led_name,))
-        led_count = led_helper.led_count
+        led_count = led_helper.get_led_count()
         index = gcmd.get_int("INDEX", None, minval=1, maxval=led_count)
         template = None
         lparams = {}