From 3aa34deaa2ad1025a8c14a615f0c2e18df0efede Mon Sep 17 00:00:00 2001
From: Dan Smith <dsmith@danplanet.com>
Date: Thu, 11 Feb 2016 17:10:34 +0000
Subject: [PATCH] Add state_as_number() helper

This adds state_as_number(), a helper method that tries to interpret
state as a number, for cases we can predict. It's a generalization of
what is copy-and-paste-ed into multiple other places.
---
 homeassistant/helpers/state.py | 34 +++++++++++++++++++++++++++--
 tests/helpers/test_state.py    | 40 ++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 2 deletions(-)

diff --git a/homeassistant/helpers/state.py b/homeassistant/helpers/state.py
index 9c74d59844e..1a3520d4733 100644
--- a/homeassistant/helpers/state.py
+++ b/homeassistant/helpers/state.py
@@ -8,8 +8,11 @@ import homeassistant.util.dt as dt_util
 from homeassistant.const import (
     STATE_ON, STATE_OFF, SERVICE_TURN_ON, SERVICE_TURN_OFF,
     SERVICE_MEDIA_PLAY, SERVICE_MEDIA_PAUSE,
-    STATE_PLAYING, STATE_PAUSED, ATTR_ENTITY_ID)
-
+    STATE_PLAYING, STATE_PAUSED, ATTR_ENTITY_ID,
+    STATE_LOCKED, STATE_UNLOCKED, STATE_UNKNOWN,
+    STATE_OPEN, STATE_CLOSED)
+from homeassistant.components.sun import (STATE_ABOVE_HORIZON,
+                                          STATE_BELOW_HORIZON)
 from homeassistant.components.media_player import (SERVICE_PLAY_MEDIA)
 
 _LOGGER = logging.getLogger(__name__)
@@ -92,3 +95,30 @@ def reproduce_state(hass, states, blocking=False):
         data = json.loads(service_data)
         data[ATTR_ENTITY_ID] = entity_ids
         hass.services.call(service_domain, service, data, blocking)
+
+
+def state_as_number(state):
+    """Try to coerce our state to a number.
+
+    Raises ValueError if this is not possible.
+    """
+
+    if state.state in (STATE_ON, STATE_LOCKED, STATE_ABOVE_HORIZON,
+                       STATE_OPEN):
+        return 1
+    elif state.state in (STATE_OFF, STATE_UNLOCKED, STATE_UNKNOWN,
+                         STATE_BELOW_HORIZON, STATE_CLOSED):
+        return 0
+    else:
+        try:
+            # This distinction is probably not important,
+            # but in case something downstream cares about
+            # int vs. float, try to be helpful here.
+            if '.' in state.state:
+                return float(state.state)
+            else:
+                return int(state.state)
+        except (ValueError, TypeError):
+            pass
+
+    raise ValueError('State is not a number')
diff --git a/tests/helpers/test_state.py b/tests/helpers/test_state.py
index f4e28330f7a..e222e72fe4b 100644
--- a/tests/helpers/test_state.py
+++ b/tests/helpers/test_state.py
@@ -13,6 +13,12 @@ import homeassistant.components as core_components
 from homeassistant.const import SERVICE_TURN_ON
 from homeassistant.util import dt as dt_util
 from homeassistant.helpers import state
+from homeassistant.const import (
+    STATE_OFF, STATE_OPEN, STATE_CLOSED,
+    STATE_LOCKED, STATE_UNLOCKED, STATE_UNKNOWN,
+    STATE_ON, STATE_OFF)
+from homeassistant.components.sun import (STATE_ABOVE_HORIZON,
+                                          STATE_BELOW_HORIZON)
 
 from tests.common import get_test_home_assistant, mock_service
 
@@ -146,3 +152,37 @@ class TestStateHelpers(unittest.TestCase):
         self.assertEqual(['light.test1', 'light.test2'],
                          last_call.data.get('entity_id'))
         self.assertEqual(95, last_call.data.get('brightness'))
+
+    def test_as_number_states(self):
+        zero_states = (STATE_OFF, STATE_CLOSED, STATE_UNLOCKED,
+                       STATE_BELOW_HORIZON)
+        one_states = (STATE_ON, STATE_OPEN, STATE_LOCKED, STATE_ABOVE_HORIZON)
+        for _state in zero_states:
+            self.assertEqual(0, state.state_as_number(
+                ha.State('domain.test', _state, {})))
+        for _state in one_states:
+            self.assertEqual(1, state.state_as_number(
+                ha.State('domain.test', _state, {})))
+
+    def test_as_number_coercion(self):
+        for _state in ('0', '0.0'):
+            self.assertEqual(
+                0.0, float(state.state_as_number(
+                ha.State('domain.test', _state, {}))))
+        for _state in ('1', '1.0'):
+            self.assertEqual(
+                1.0, float(state.state_as_number(
+                    ha.State('domain.test', _state, {}))))
+
+    def test_as_number_tries_to_keep_types(self):
+        result = state.state_as_number(ha.State('domain.test', '1', {}))
+        self.assertTrue(isinstance(result, int))
+        result = state.state_as_number(ha.State('domain.test', '1.0', {}))
+        self.assertTrue(isinstance(result, float))
+
+    def test_as_number_invalid_cases(self):
+        for _state in ('', 'foo', 'foo.bar', None, False, True, None,
+                       object, object()):
+            self.assertRaises(ValueError,
+                              state.state_as_number,
+                              ha.State('domain.test', _state, {}))