From 92b6ed41809e947b5a6fd9a737daebe55f478de7 Mon Sep 17 00:00:00 2001
From: Otto Winter <otto@otto-winter.com>
Date: Sat, 20 Oct 2018 13:14:02 +0200
Subject: [PATCH] Add FastLED color correction option (#200)

* Add FastLED color correction option

* Add test
---
 esphomeyaml/components/light/fastled_clockless.py | 7 ++++++-
 esphomeyaml/components/light/fastled_spi.py       | 9 +++++++--
 esphomeyaml/const.py                              | 1 +
 tests/test1.yaml                                  | 1 +
 4 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/esphomeyaml/components/light/fastled_clockless.py b/esphomeyaml/components/light/fastled_clockless.py
index 1326853bd5..767703c152 100644
--- a/esphomeyaml/components/light/fastled_clockless.py
+++ b/esphomeyaml/components/light/fastled_clockless.py
@@ -6,7 +6,7 @@ from esphomeyaml.components import light
 from esphomeyaml.components.power_supply import PowerSupplyComponent
 from esphomeyaml.const import CONF_CHIPSET, CONF_DEFAULT_TRANSITION_LENGTH, CONF_GAMMA_CORRECT, \
     CONF_MAKE_ID, CONF_MAX_REFRESH_RATE, CONF_NAME, CONF_NUM_LEDS, CONF_PIN, CONF_POWER_SUPPLY, \
-    CONF_RGB_ORDER, CONF_EFFECTS
+    CONF_RGB_ORDER, CONF_EFFECTS, CONF_COLOR_CORRECT
 from esphomeyaml.helpers import App, Application, RawExpression, TemplateArguments, add, \
     get_variable, variable
 
@@ -66,6 +66,7 @@ PLATFORM_SCHEMA = cv.nameable(light.LIGHT_PLATFORM_SCHEMA.extend({
     vol.Optional(CONF_RGB_ORDER): vol.All(vol.Upper, cv.one_of(*RGB_ORDERS)),
 
     vol.Optional(CONF_GAMMA_CORRECT): cv.positive_float,
+    vol.Optional(CONF_COLOR_CORRECT): vol.All([cv.percentage], vol.Length(min=3, max=3)),
     vol.Optional(CONF_DEFAULT_TRANSITION_LENGTH): cv.positive_time_period_milliseconds,
     vol.Optional(CONF_POWER_SUPPLY): cv.use_variable_id(PowerSupplyComponent),
     vol.Optional(CONF_EFFECTS): light.validate_effects(light.FASTLED_EFFECTS),
@@ -93,6 +94,10 @@ def to_code(config):
             yield
         add(fast_led.set_power_supply(power_supply))
 
+    if CONF_COLOR_CORRECT in config:
+        r, g, b = config[CONF_COLOR_CORRECT]
+        add(fast_led.set_correction(r, g, b))
+
     light.setup_light(make.Pstate, make.Pmqtt, config)
 
 
diff --git a/esphomeyaml/components/light/fastled_spi.py b/esphomeyaml/components/light/fastled_spi.py
index ec1ccb3509..a74a435d5b 100644
--- a/esphomeyaml/components/light/fastled_spi.py
+++ b/esphomeyaml/components/light/fastled_spi.py
@@ -6,7 +6,7 @@ from esphomeyaml.components import light
 from esphomeyaml.components.power_supply import PowerSupplyComponent
 from esphomeyaml.const import CONF_CHIPSET, CONF_CLOCK_PIN, CONF_DATA_PIN, \
     CONF_DEFAULT_TRANSITION_LENGTH, CONF_GAMMA_CORRECT, CONF_MAKE_ID, CONF_MAX_REFRESH_RATE, \
-    CONF_NAME, CONF_NUM_LEDS, CONF_POWER_SUPPLY, CONF_RGB_ORDER, CONF_EFFECTS
+    CONF_NAME, CONF_NUM_LEDS, CONF_POWER_SUPPLY, CONF_RGB_ORDER, CONF_EFFECTS, CONF_COLOR_CORRECT
 from esphomeyaml.helpers import App, Application, RawExpression, TemplateArguments, add, \
     get_variable, variable
 
@@ -44,6 +44,7 @@ PLATFORM_SCHEMA = cv.nameable(light.LIGHT_PLATFORM_SCHEMA.extend({
     vol.Optional(CONF_MAX_REFRESH_RATE): cv.positive_time_period_microseconds,
 
     vol.Optional(CONF_GAMMA_CORRECT): cv.positive_float,
+    vol.Optional(CONF_COLOR_CORRECT): vol.All([cv.percentage], vol.Length(min=3, max=3)),
     vol.Optional(CONF_DEFAULT_TRANSITION_LENGTH): cv.positive_time_period_milliseconds,
     vol.Optional(CONF_POWER_SUPPLY): cv.use_variable_id(PowerSupplyComponent),
     vol.Optional(CONF_EFFECTS): light.validate_effects(light.FASTLED_EFFECTS),
@@ -73,7 +74,11 @@ def to_code(config):
             yield
         add(fast_led.set_power_supply(power_supply))
 
-        light.setup_light(make.Pstate, make.Pmqtt, config)
+    if CONF_COLOR_CORRECT in config:
+        r, g, b = config[CONF_COLOR_CORRECT]
+        add(fast_led.set_correction(r, g, b))
+
+    light.setup_light(make.Pstate, make.Pmqtt, config)
 
 
 BUILD_FLAGS = '-DUSE_FAST_LED_LIGHT'
diff --git a/esphomeyaml/const.py b/esphomeyaml/const.py
index 2c89d8cc80..f67061ef24 100644
--- a/esphomeyaml/const.py
+++ b/esphomeyaml/const.py
@@ -346,6 +346,7 @@ CONF_PM_2_5 = 'pm_2_5'
 CONF_PM_10_0 = 'pm_10_0'
 CONF_FORMALDEHYDE = 'formaldehyde'
 CONF_ON_TAG = 'on_tag'
+CONF_COLOR_CORRECT = 'color_correct'
 CONF_ON_JSON_MESSAGE = 'on_json_message'
 
 ALLOWED_NAME_CHARS = u'abcdefghijklmnopqrstuvwxyz0123456789_'
diff --git a/tests/test1.yaml b/tests/test1.yaml
index fc44caa29b..deb741dfeb 100644
--- a/tests/test1.yaml
+++ b/tests/test1.yaml
@@ -649,6 +649,7 @@ light:
     rgb_order: BRG
     max_refresh_rate: 20ms
     power_supply: atx_power_supply
+    color_correct: [75%, 100%, 50%]
     name: "FastLED WS2811 Light"
     effects:
     - fastled_color_wipe: