Adding color support to template light

homeassistant/components/light/template.py

@@ -11,7 +11,7 @@
 
 from homeassistant.core import callback
 from homeassistant.components.light import (
-    ATTR_BRIGHTNESS, ENTITY_ID_FORMAT, Light, SUPPORT_BRIGHTNESS)
+    ATTR_BRIGHTNESS, ATTR_RGB_COLOR, ENTITY_ID_FORMAT, Light, SUPPORT_BRIGHTNESS, SUPPORT_RGB_COLOR)
 from homeassistant.const import (
     CONF_VALUE_TEMPLATE, CONF_ICON_TEMPLATE, CONF_ENTITY_PICTURE_TEMPLATE,
     CONF_ENTITY_ID, CONF_FRIENDLY_NAME, STATE_ON, STATE_OFF,
@@ -31,6 +31,8 @@
 CONF_OFF_ACTION = 'turn_off'
 CONF_LEVEL_ACTION = 'set_level'
 CONF_LEVEL_TEMPLATE = 'level_template'
+CONF_COLOR_ACTION = 'set_color'
+CONF_COLOR_TEMPLATE = 'color_template'
 
 LIGHT_SCHEMA = vol.Schema({
     vol.Required(CONF_ON_ACTION): cv.SCRIPT_SCHEMA,
@@ -40,10 +42,17 @@
     vol.Optional(CONF_ENTITY_PICTURE_TEMPLATE, default=None): cv.template,
     vol.Optional(CONF_LEVEL_ACTION, default=None): cv.SCRIPT_SCHEMA,
     vol.Optional(CONF_LEVEL_TEMPLATE, default=None): cv.template,
+    vol.Optional(CONF_COLOR_ACTION, default=None): cv.SCRIPT_SCHEMA,
+    vol.Optional(CONF_COLOR_TEMPLATE, default=None): cv.template,
     vol.Optional(CONF_FRIENDLY_NAME): cv.string,
     vol.Optional(CONF_ENTITY_ID): cv.entity_ids
 })
 
+LIGHT_SCHEMA = vol.All(
+    cv.deprecated(CONF_ENTITY_ID),
+    LIGHT_SCHEMA,
+)
+
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
     vol.Required(CONF_LIGHTS): vol.Schema({cv.slug: LIGHT_SCHEMA}),
 })
@@ -64,6 +73,8 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
         off_action = device_config[CONF_OFF_ACTION]
         level_action = device_config.get(CONF_LEVEL_ACTION)
         level_template = device_config[CONF_LEVEL_TEMPLATE]
+        color_action = device_config.get(CONF_COLOR_ACTION)
+        color_template = device_config[CONF_COLOR_TEMPLATE]
 
         template_entity_ids = set()
 
@@ -77,6 +88,11 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
             if str(temp_ids) != MATCH_ALL:
                 template_entity_ids |= set(temp_ids)
 
+        if color_template is not None:
+            temp_ids = color_template.extract_entities()
+            if str(temp_ids) != MATCH_ALL:
+                template_entity_ids |= set(temp_ids)
+
         if icon_template is not None:
             temp_ids = icon_template.extract_entities()
             if str(temp_ids) != MATCH_ALL:
@@ -96,7 +112,7 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
             LightTemplate(
                 hass, device, friendly_name, state_template,
                 icon_template, entity_picture_template, on_action,
-                off_action, level_action, level_template, entity_ids)
+                off_action, level_action, level_template, color_action, color_template, entity_ids)
         )
 
     if not lights:
@@ -108,11 +124,11 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
 
 
 class LightTemplate(Light):
-    """Representation of a templated Light, including dimmable."""
+    """Representation of a templated Light, including dimmable and color."""
 
     def __init__(self, hass, device_id, friendly_name, state_template,
                  icon_template, entity_picture_template, on_action,
-                 off_action, level_action, level_template, entity_ids):
+                 off_action, level_action, level_template, color_action, color_template, entity_ids):
         """Initialize the light."""
         self.hass = hass
         self.entity_id = async_generate_entity_id(
@@ -127,17 +143,24 @@ def __init__(self, hass, device_id, friendly_name, state_template,
         if level_action is not None:
             self._level_script = Script(hass, level_action)
         self._level_template = level_template
+        self._color_script = None
+        if color_action is not None:
+            self._color_script = Script(hass, color_action)
+        self._color_template = color_template
 
         self._state = False
         self._icon = None
         self._entity_picture = None
         self._brightness = None
+        self._color = None
         self._entities = entity_ids
 
         if self._template is not None:
             self._template.hass = self.hass
         if self._level_template is not None:
             self._level_template.hass = self.hass
+        if self._color_template is not None:
+            self._color_template.hass = self.hass
         if self._icon_template is not None:
             self._icon_template.hass = self.hass
         if self._entity_picture_template is not None:
@@ -148,6 +171,11 @@ def brightness(self):
         """Return the brightness of the light."""
         return self._brightness
 
+    @property
+    def color(self):
+        """Return the color of the light."""
+        return self._color
+
     @property
     def name(self):
         """Return the display name of this light."""
@@ -156,9 +184,12 @@ def name(self):
     @property
     def supported_features(self):
         """Flag supported features."""
-        if self._level_script is not None:
+        if self._level_script is not None and self._color_script is not None:
+            return SUPPORT_RGB_COLOR |SUPPORT_BRIGHTNESS
+        elif self._level_script is not None:
             return SUPPORT_BRIGHTNESS
-
+        elif self._color_script is not None:
+            return SUPPORT_RGB_COLOR
         return 0
 
     @property
@@ -193,7 +224,8 @@ def template_light_state_listener(entity, old_state, new_state):
         def template_light_startup(event):
             """Update template on startup."""
             if (self._template is not None or
-                    self._level_template is not None):
+                    self._level_template is not None or
+                    self._color_template is not None):
                 async_track_state_change(
                     self.hass, self._entities, template_light_state_listener)
 
@@ -217,10 +249,22 @@ def async_turn_on(self, **kwargs):
             self._brightness = kwargs[ATTR_BRIGHTNESS]
             optimistic_set = True
 
+        if self._color_template is None and ATTR_RGB_COLOR in kwargs:
+            _LOGGER.info("Optimistically setting color to %s",
+                         kwargs[ATTR_RGB_COLOR])
+            self._color = kwargs[ATTR_RGB_COLOR]
+            optimistic_set = True
+
         if ATTR_BRIGHTNESS in kwargs and self._level_script:
             self.hass.async_add_job(self._level_script.async_run(
                 {"brightness": kwargs[ATTR_BRIGHTNESS]}))
-        else:
+
+        if ATTR_RGB_COLOR in kwargs and self._color_script:
+            self.hass.async_add_job(self._color_script.async_run(
+                {"color": kwargs[ATTR_RGB_COLOR]}))
+
+        if not (ATTR_BRIGHTNESS in kwargs and self._level_script and
+                ATTR_RGB_COLOR in kwargs and self._color_script):
             yield from self._on_script.async_run()
 
         if optimistic_set:
@@ -270,6 +314,16 @@ def async_update(self):
                     brightness)
                 self._brightness = None
 
+        if self._color_template is not None:
+            try:
+                color = self._color_template.async_render()
+            except TemplateError as ex:
+                _LOGGER.error(ex)
+                self._state = None
+
+            self._color = color
+            #TODO input check for colors
+
         for property_name, template in (
                 ('_icon', self._icon_template),
                 ('_entity_picture', self._entity_picture_template)):