From f73f824e0e5df55bb9b86c6771b26c0d262e1d95 Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <paulus@paulusschoutsen.nl>
Date: Sat, 12 Dec 2015 22:19:37 -0800
Subject: [PATCH] Make template states var callable

---
 homeassistant/util/template.py | 14 +++++++++++++-
 tests/util/test_template.py    | 17 +++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/homeassistant/util/template.py b/homeassistant/util/template.py
index f8fa2c70f8bd..d0a07507bdf0 100644
--- a/homeassistant/util/template.py
+++ b/homeassistant/util/template.py
@@ -9,6 +9,7 @@ import json
 import logging
 import jinja2
 from jinja2.sandbox import ImmutableSandboxedEnvironment
+from homeassistant.const import STATE_UNKNOWN
 from homeassistant.exceptions import TemplateError
 
 _LOGGER = logging.getLogger(__name__)
@@ -60,6 +61,10 @@ class AllStates(object):
         return iter(sorted(self._hass.states.all(),
                            key=lambda state: state.entity_id))
 
+    def __call__(self, entity_id):
+        state = self._hass.states.get(entity_id)
+        return STATE_UNKNOWN if state is None else state.state
+
 
 class DomainStates(object):
     """ Class to expose a specific HA domain as attributes. """
@@ -96,6 +101,13 @@ def multiply(value, amount):
         # If value can't be converted to float
         return value
 
-ENV = ImmutableSandboxedEnvironment()
+
+class TemplateEnvironment(ImmutableSandboxedEnvironment):
+    """ Home Assistant template environment. """
+
+    def is_safe_callable(self, obj):
+        return isinstance(obj, AllStates) or super().is_safe_callable(obj)
+
+ENV = TemplateEnvironment()
 ENV.filters['round'] = forgiving_round
 ENV.filters['multiply'] = multiply
diff --git a/tests/util/test_template.py b/tests/util/test_template.py
index 1e34d999fa75..1ecd7d5b894e 100644
--- a/tests/util/test_template.py
+++ b/tests/util/test_template.py
@@ -101,6 +101,14 @@ class TestUtilTemplate(unittest.TestCase):
         with self.assertRaises(TemplateError):
             template.render(self.hass, '{{ invalid_syntax')
 
+    def test_if_state_exists(self):
+        self.hass.states.set('test.object', 'available')
+        self.assertEqual(
+            'exists',
+            template.render(
+                self.hass,
+                '{% if states.test.object %}exists{% else %}not exists{% endif %}'))
+
     def test_is_state(self):
         self.hass.states.set('test.object', 'available')
         self.assertEqual(
@@ -108,3 +116,12 @@ class TestUtilTemplate(unittest.TestCase):
             template.render(
                 self.hass,
                 '{% if is_state("test.object", "available") %}yes{% else %}no{% endif %}'))
+
+    def test_states_function(self):
+        self.hass.states.set('test.object', 'available')
+        self.assertEqual(
+            'available',
+            template.render(self.hass, '{{ states("test.object") }}'))
+        self.assertEqual(
+            'unknown',
+            template.render(self.hass, '{{ states("test.object2") }}'))