Improve Python script

homeassistant/components/python_script.py

@@ -2,6 +2,7 @@
 import glob
 import os
 import logging
+import datetime
 
 import voluptuous as vol
 
@@ -63,7 +64,8 @@ def execute_script(hass, name, data=None):
 def execute(hass, filename, source, data=None):
     """Execute Python source."""
     from RestrictedPython import compile_restricted_exec
-    from RestrictedPython.Guards import safe_builtins, full_write_guard
+    from RestrictedPython.Guards import safe_builtins, full_write_guard, \
+        guarded_iter_unpack_sequence, guarded_unpack_sequence
     from RestrictedPython.Utilities import utility_builtins
     from RestrictedPython.Eval import default_guarded_getitem
 
@@ -94,13 +96,16 @@ def protected_getattr(obj, name, default=None):
 
     builtins = safe_builtins.copy()
     builtins.update(utility_builtins)
+    builtins['datetime'] = datetime
     restricted_globals = {
         '__builtins__': builtins,
         '_print_': StubPrinter,
         '_getattr_': protected_getattr,
         '_write_': full_write_guard,
         '_getiter_': iter,
-        '_getitem_': default_guarded_getitem
+        '_getitem_': default_guarded_getitem,
+        '_iter_unpack_sequence_': guarded_iter_unpack_sequence,
+        '_unpack_sequence_': guarded_unpack_sequence,
     }
     logger = logging.getLogger('{}.{}'.format(__name__, filename))
     local = {

tests/components/test_python_script.py

@@ -180,3 +180,26 @@ def test_iterating(hass):
 
     assert hass.states.is_state('hello.1', 'world')
     assert hass.states.is_state('hello.2', 'world')
+
+
+@asyncio.coroutine
+def test_unpacking_sequence(hass, caplog):
+    """Test compile error logs error."""
+    caplog.set_level(logging.ERROR)
+    source = """
+a,b = (1,2)
+ab_list = [(a,b) for a,b in [(1, 2), (3, 4)]]
+hass.states.set('hello.a', a)
+hass.states.set('hello.b', b)
+hass.states.set('hello.ab_list', '{}'.format(ab_list))
+"""
+
+    hass.async_add_job(execute, hass, 'test.py', source, {})
+    yield from hass.async_block_till_done()
+
+    assert hass.states.is_state('hello.a', '1')
+    assert hass.states.is_state('hello.b', '2')
+    assert hass.states.is_state('hello.ab_list', '[(1, 2), (3, 4)]')
+
+    # No errors logged = good
+    assert caplog.text == ''