Add state_with_unit property to state objects in templates (#9014)

* Wrap state objects in templates

* Fix tests

* Fix bugs

* Lint

* Remove invalid state warning

homeassistant/helpers/template.py

@@ -10,7 +10,8 @@ from jinja2 import contextfilter
 from jinja2.sandbox import ImmutableSandboxedEnvironment
 
 from homeassistant.const import (
-    STATE_UNKNOWN, ATTR_LATITUDE, ATTR_LONGITUDE, MATCH_ALL)
+    STATE_UNKNOWN, ATTR_LATITUDE, ATTR_LONGITUDE, MATCH_ALL,
+    ATTR_UNIT_OF_MEASUREMENT)
 from homeassistant.core import State
 from homeassistant.exceptions import TemplateError
 from homeassistant.helpers import location as loc_helper
@@ -181,8 +182,10 @@ class AllStates(object):
 
     def __iter__(self):
         """Return all states."""
-        return iter(sorted(self._hass.states.async_all(),
-                           key=lambda state: state.entity_id))
+        return iter(
+            _wrap_state(state) for state in
+            sorted(self._hass.states.async_all(),
+                   key=lambda state: state.entity_id))
 
     def __call__(self, entity_id):
         """Return the states."""
@@ -200,7 +203,8 @@ class DomainStates(object):
 
     def __getattr__(self, name):
         """Return the states."""
-        return self._hass.states.get('{}.{}'.format(self._domain, name))
+        return _wrap_state(
+            self._hass.states.get('{}.{}'.format(self._domain, name)))
 
     def __iter__(self):
         """Return the iteration over all the states."""
@@ -210,6 +214,42 @@ class DomainStates(object):
             key=lambda state: state.entity_id))
 
 
+class TemplateState(State):
+    """Class to represent a state object in a template."""
+
+    # Inheritance is done so functions that check against State keep working
+    # pylint: disable=super-init-not-called
+    def __init__(self, state):
+        """Initialize template state."""
+        self._state = state
+
+    @property
+    def state_with_unit(self):
+        """Return the state concatenated with the unit if available."""
+        state = object.__getattribute__(self, '_state')
+        unit = state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
+        if unit is None:
+            return state.state
+        return "{} {}".format(state.state, unit)
+
+    def __getattribute__(self, name):
+        """Return an attribute of the state."""
+        if name in TemplateState.__dict__:
+            return object.__getattribute__(self, name)
+        else:
+            return getattr(object.__getattribute__(self, '_state'), name)
+
+    def __repr__(self):
+        """Representation of Template State."""
+        rep = object.__getattribute__(self, '_state').__repr__()
+        return '<template ' + rep[1:]
+
+
+def _wrap_state(state):
+    """Helper function to wrap a state."""
+    return None if state is None else TemplateState(state)
+
+
 class LocationMethods(object):
     """Class to expose distance helpers to templates."""
 
@@ -278,7 +318,7 @@ class LocationMethods(object):
             states = [self._hass.states.get(entity_id) for entity_id
                       in group.expand_entity_ids(self._hass, [gr_entity_id])]
 
-        return loc_helper.closest(latitude, longitude, states)
+        return _wrap_state(loc_helper.closest(latitude, longitude, states))
 
     def distance(self, *args):
         """Calculate distance.

tests/helpers/test_template.py

@@ -1,4 +1,5 @@
 """Test Home Assistant template helper methods."""
+import asyncio
 from datetime import datetime
 import unittest
 import random
@@ -652,8 +653,9 @@ class TestHelpersTemplate(unittest.TestCase):
     def test_closest_function_no_location_states(self):
         """Test closest function without location states."""
         self.assertEqual(
-            'None',
-            template.Template('{{ closest(states) }}', self.hass).render())
+            '',
+            template.Template('{{ closest(states).entity_id }}',
+                              self.hass).render())
 
     def test_extract_entities_none_exclude_stuff(self):
         """Test extract entities function with none or exclude stuff."""
@@ -750,3 +752,33 @@ is_state_attr('device_tracker.phone_2', 'battery', 40)
                 " > (states('input_slider.luftfeuchtigkeit') | int +1.5)"
                 " %}true{% endif %}"
             )))
+
+
+@asyncio.coroutine
+def test_state_with_unit(hass):
+    """Test the state_with_unit property helper."""
+    hass.states.async_set('sensor.test', '23', {
+        'unit_of_measurement': 'beers',
+    })
+    hass.states.async_set('sensor.test2', 'wow')
+
+    tpl = template.Template(
+        '{{ states.sensor.test.state_with_unit }}', hass)
+
+    assert tpl.async_render() == '23 beers'
+
+    tpl = template.Template(
+        '{{ states.sensor.test2.state_with_unit }}', hass)
+
+    assert tpl.async_render() == 'wow'
+
+    tpl = template.Template(
+        '{% for state in states %}{{ state.state_with_unit }} {% endfor %}',
+        hass)
+
+    assert tpl.async_render() == '23 beers wow'
+
+    tpl = template.Template('{{ states.sensor.non_existing.state_with_unit }}',
+                            hass)
+
+    assert tpl.async_render() == ''