diff --git a/homeassistant/helpers/template.py b/homeassistant/helpers/template.py index 6c74c49424e..aa6ca186a8e 100644 --- a/homeassistant/helpers/template.py +++ b/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 '