update led payload options (#61)

rae_sdk/rae_sdk/robot/led.py

@@ -14,18 +14,18 @@ class LEDController:
     Methods
     -------
         set_leds(payload): Sets the robot's LEDs to a given color.
-    
+
     """
-    
+
     def __init__(self, ros_interface):
         self._ros_interface = ros_interface
         self._ros_interface.create_publisher("/leds", LEDControl)
         self._effect_types_map = {
-            "all" : LEDControl.CTRL_TYPE_ALL,
-            "single" : LEDControl.CTRL_TYPE_SINGLE,
-            "spinner" : LEDControl.CTRL_TYPE_SPINNER,
-            "fan" : LEDControl.CTRL_TYPE_FAN,
-            "custom" : LEDControl.CTRL_TYPE_CUSTOM,
+            "all": LEDControl.CTRL_TYPE_ALL,
+            "single": LEDControl.CTRL_TYPE_SINGLE,
+            "spinner": LEDControl.CTRL_TYPE_SPINNER,
+            "fan": LEDControl.CTRL_TYPE_FAN,
+            "custom": LEDControl.CTRL_TYPE_CUSTOM,
         }
         log.info("LED Controller ready")
 
@@ -42,28 +42,37 @@ def normalize(self, num):
         else:
             return float(num)/255.0
 
-    def set_leds_from_payload(self, payload):
+    def set_leds_from_payload(self, payload: dict):
         """
         Set the robot's LEDs to a given color.
-        
+
         Args:
         ----
             payload (dict): A dictionary containing the color to set the LEDs to. 
             Example payload struct: {'brightness': 50, 'color': '#FFFFFF', 'effect': 'pulse', 'interval': 5}
-            
+
         """
         led_msg = LEDControl()
 
         r, g, b = self.hex_to_rgb(payload['color'])
-
         color_msg = ColorPeriod()
-        color_msg.period = payload['interval']
+        if 'interval' in payload.keys():
+            color_msg.frequency = float(payload['interval'])
+        else:
+            color_msg.frequency = 1.0
         color_msg.color.a = 1.0
         color_msg.color.r = self.normalize(r)
         color_msg.color.g = self.normalize(g)
         color_msg.color.b = self.normalize(b)
         led_msg.data = [color_msg]
-        led_msg.control_type = led_msg.CTRL_TYPE_ALL
+        if payload['effect'] in self._effect_types_map:
+            led_msg.control_type = self._effect_types_map[payload['effect']]
+        else:
+            led_msg.control_type = LEDControl.CTRL_TYPE_ALL
+        if 'size' in payload.keys():
+            led_msg.animation_size = payload['size']
+        if 'blades' in payload.keys():
+            led_msg.animation_quantity = payload['blades']
         self._ros_interface.publish("/leds", led_msg)
 
     def set_leds(self, color: str, brightness: int = 100, effect: str = "solid", interval: int = 5):
@@ -76,7 +85,7 @@ def set_leds(self, color: str, brightness: int = 100, effect: str = "solid", int
             brightness (int): The brightness of the LEDs. (Default: 100)
             effect (str): The effect to apply to the LEDs. (Default: "solid")
             interval (int): The interval of the effect. (Default: 5)
-            
+
         """
         payload = {
             'brightness': brightness,
@@ -85,14 +94,14 @@ def set_leds(self, color: str, brightness: int = 100, effect: str = "solid", int
             'interval': interval
         }
         self.set_leds_from_payload(payload)
-    
+
     def set_leds_from_msg(self, msg: LEDControl):
         """
         Set the robot's LEDs to a given color.
 
         Args:
         ----
             msg (LEDControl): The message containing the color to set the LEDs to.
-            
+
         """
-        self._ros_interface.publish("/leds", msg)
+        self._ros_interface.publish("/leds", msg)