From 8d1999dc1205fb8c2156b0a3ed66f1320574b7ed Mon Sep 17 00:00:00 2001
From: Sean Dague <sean@dague.net>
Date: Thu, 20 Jul 2017 01:56:24 -0400
Subject: [PATCH] Enhance python_script to support "_getitem_" (#8541)

* Enhance python_script to support "_getitem_"

In order to use dict / list structures in python scripts we need
_getitem_ allowed in the RestrictedPython environment. There is a
default_guarded_getitem included with RestrictedPython, which is a
pass through used in the Eval code paths.

* Add tests for dict/list support in python_scripts

* Lint
---
 homeassistant/components/python_script.py |  2 ++
 tests/components/test_python_script.py    | 16 ++++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/homeassistant/components/python_script.py b/homeassistant/components/python_script.py
index 66b9bd0112a..386abba59ae 100644
--- a/homeassistant/components/python_script.py
+++ b/homeassistant/components/python_script.py
@@ -65,6 +65,7 @@ def execute(hass, filename, source, data=None):
     from RestrictedPython import compile_restricted_exec
     from RestrictedPython.Guards import safe_builtins, full_write_guard
     from RestrictedPython.Utilities import utility_builtins
+    from RestrictedPython.Eval import default_guarded_getitem
 
     compiled = compile_restricted_exec(source, filename=filename)
 
@@ -99,6 +100,7 @@ def execute(hass, filename, source, data=None):
         '_getattr_': protected_getattr,
         '_write_': full_write_guard,
         '_getiter_': iter,
+        '_getitem_': default_guarded_getitem
     }
     logger = logging.getLogger('{}.{}'.format(__name__, filename))
     local = {
diff --git a/tests/components/test_python_script.py b/tests/components/test_python_script.py
index 78142e06cf3..62c1b67eba9 100644
--- a/tests/components/test_python_script.py
+++ b/tests/components/test_python_script.py
@@ -137,6 +137,22 @@ hass.async_stop()
     assert "Not allowed to access async methods" in caplog.text
 
 
+@asyncio.coroutine
+def test_using_complex_structures(hass, caplog):
+    """Test that dicts and lists work."""
+    caplog.set_level(logging.INFO)
+    source = """
+mydict = {"a": 1, "b": 2}
+mylist = [1, 2, 3, 4]
+logger.info('Logging from inside script: %s %s' % (mydict["a"], mylist[2]))
+    """
+
+    hass.async_add_job(execute, hass, 'test.py', source, {})
+    yield from hass.async_block_till_done()
+
+    assert "Logging from inside script: 1 3" in caplog.text
+
+
 @asyncio.coroutine
 def test_accessing_forbidden_methods(hass, caplog):
     """Test compile error logs error."""