Enhance python_script to support "_getitem_"

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 protected_getattr(obj, name, default=None):
         '_getattr_': protected_getattr,
         '_write_': full_write_guard,
         '_getiter_': iter,
+        '_getitem_': default_guarded_getitem
     }
     logger = logging.getLogger('{}.{}'.format(__name__, filename))
     local = {

tests/components/test_python_script.py

@@ -137,6 +137,22 @@ def test_accessing_async_methods(hass, caplog):
     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."""