From 8c010c8df434f63680579db6fcbeba20483159ed Mon Sep 17 00:00:00 2001
From: Philip Lundrigan <philiplundrigan@gmail.com>
Date: Mon, 21 Dec 2015 16:09:27 -0700
Subject: [PATCH] Add ability to use sun as condition in automation

---
 homeassistant/components/automation/sun.py | 114 +++++++++++++++++----
 1 file changed, 94 insertions(+), 20 deletions(-)

diff --git a/homeassistant/components/automation/sun.py b/homeassistant/components/automation/sun.py
index 84334493d0f..cf14d220fb0 100644
--- a/homeassistant/components/automation/sun.py
+++ b/homeassistant/components/automation/sun.py
@@ -17,6 +17,10 @@ DEPENDENCIES = ['sun']
 
 CONF_OFFSET = 'offset'
 CONF_EVENT = 'event'
+CONF_BEFORE = "before"
+CONF_BEFORE_OFFSET = "before_offset"
+CONF_AFTER = "after"
+CONF_AFTER_OFFSET = "after_offset"
 
 EVENT_SUNSET = 'sunset'
 EVENT_SUNRISE = 'sunrise'
@@ -37,26 +41,9 @@ def trigger(hass, config, action):
         _LOGGER.error("Invalid value for %s: %s", CONF_EVENT, event)
         return False
 
-    if CONF_OFFSET in config:
-        raw_offset = config.get(CONF_OFFSET)
-
-        negative_offset = False
-        if raw_offset.startswith('-'):
-            negative_offset = True
-            raw_offset = raw_offset[1:]
-
-        try:
-            (hour, minute, second) = [int(x) for x in raw_offset.split(':')]
-        except ValueError:
-            _LOGGER.error('Could not parse offset %s', raw_offset)
-            return False
-
-        offset = timedelta(hours=hour, minutes=minute, seconds=second)
-
-        if negative_offset:
-            offset *= -1
-    else:
-        offset = timedelta(0)
+    offset = _parse_offset(config.get(CONF_OFFSET))
+    if offset is False:
+        return False
 
     # Do something to call action
     if event == EVENT_SUNRISE:
@@ -67,6 +54,70 @@ def trigger(hass, config, action):
     return True
 
 
+def if_action(hass, config):
+    """ Wraps action method with sun based condition. """
+    before = config.get(CONF_BEFORE)
+    after = config.get(CONF_AFTER)
+
+    # Make sure required configuration keys are present
+    if before is None and after is None:
+        logging.getLogger(__name__).error(
+            "Missing if-condition configuration key %s or %s",
+            CONF_BEFORE, CONF_AFTER)
+        return None
+
+    # Make sure configuration keys have the right value
+    if before is not None and before not in (EVENT_SUNRISE, EVENT_SUNSET) or \
+       after is not None and after not in (EVENT_SUNRISE, EVENT_SUNSET):
+        logging.getLogger(__name__).error(
+            "%s and %s can only be set to %s or %s",
+            CONF_BEFORE, CONF_AFTER, EVENT_SUNRISE, EVENT_SUNSET)
+        return None
+
+    before_offset = _parse_offset(config.get(CONF_BEFORE_OFFSET))
+    after_offset = _parse_offset(config.get(CONF_AFTER_OFFSET))
+    if before_offset is False or after_offset is False:
+        return None
+
+    if before is None:
+        before_func = lambda: None
+    elif before == EVENT_SUNRISE:
+        before_func = lambda: sun.next_rising_utc(hass) + before_offset
+    else:
+        before_func = lambda: sun.next_setting_utc(hass) + before_offset
+
+    if after is None:
+        after_func = lambda: None
+    elif after == EVENT_SUNRISE:
+        after_func = lambda: sun.next_rising_utc(hass) + after_offset
+    else:
+        after_func = lambda: sun.next_setting_utc(hass) + after_offset
+
+    # This is needed for testing
+    time_func = dt_util.utcnow
+
+    def time_if():
+        """ Validate time based if-condition """
+
+        # This is needed for testing.
+        nonlocal time_func
+        now = time_func()
+        before = before_func()
+        after = after_func()
+
+        if before is not None and now > now.replace(hour=before.hour,
+                                                    minute=before.minute):
+            return False
+
+        if after is not None and now < now.replace(hour=after.hour,
+                                                   minute=after.minute):
+            return False
+
+        return True
+
+    return time_if
+
+
 def trigger_sunrise(hass, action, offset):
     """ Trigger action at next sun rise. """
     def next_rise():
@@ -103,3 +154,26 @@ def trigger_sunset(hass, action, offset):
         action()
 
     track_point_in_utc_time(hass, sunset_automation_listener, next_set())
+
+
+def _parse_offset(raw_offset):
+    if raw_offset is None:
+        return timedelta(0)
+
+    negative_offset = False
+    if raw_offset.startswith('-'):
+        negative_offset = True
+        raw_offset = raw_offset[1:]
+
+    try:
+        (hour, minute, second) = [int(x) for x in raw_offset.split(':')]
+    except ValueError:
+        _LOGGER.error('Could not parse offset %s', raw_offset)
+        return False
+
+    offset = timedelta(hours=hour, minutes=minute, seconds=second)
+
+    if negative_offset:
+        offset *= -1
+
+    return offset