correctly handle iterators in tuple assignments; fixes #15

Author: craigbarratt

custom_components/pyscript/eval.py

@@ -602,17 +602,19 @@ async def recurse_assign(self, lhs, val):
         """Recursive assignment."""
         if isinstance(lhs, ast.Tuple):
             try:
-                val_len = len(val)
-            except TypeError:
-                raise TypeError("cannot unpack non-iterable object")
-            if len(lhs.elts) < val_len:
-                raise ValueError(
-                    f"too many values to unpack (expected {len(lhs.elts)})"
-                )
-            if len(lhs.elts) > val_len:
-                raise ValueError(f"too few values to unpack (expected {len(lhs.elts)})")
-            for lhs_elt, val_elt in zip(lhs.elts, val):
+                val_iter = val.__iter__()
+            except AttributeError:
+                 raise TypeError("cannot unpack non-iterable object")
+            sentinel = object()
+            for lhs_elt in lhs.elts:
+                val_elt = next(val_iter, sentinel)
+                if val_elt is sentinel:
+                    raise TypeError(
+                        f"too few values to unpack (expected {len(lhs.elts)})"
+                    )
                 await self.recurse_assign(lhs_elt, val_elt)
+            if next(val_iter, sentinel) is not sentinel:
+                 raise TypeError(f"too many values to unpack (expected {len(lhs.elts)})")
         elif isinstance(lhs, ast.Subscript):
             var = await self.aeval(lhs.value)
             if isinstance(lhs.slice, ast.Index):

tests/custom_components/pyscript/test_unit_eval.py

@@ -55,6 +55,7 @@
     ["z = [1,2,3]; [z[1], z[-1]]", [2, 3]],
     ["'{1} {0}'.format('one', 'two')", "two one"],
     ["'%d, %d' % (23, 45)", "23, 45"],
+    ["x = [[1,2,3]]; sum(*x)", 6],
     ["args = [1, 5, 10]; {6, *args, 15}", {1, 5, 6, 10, 15}],
     ["args = [1, 5, 10]; [6, *args, 15]", [6, 1, 5, 10, 15]],
     ["kw = {'x': 1, 'y': 5}; {**kw}", {"x": 1, "y": 5}],
@@ -145,6 +146,12 @@
         "z = [1,2,3]; ((x, y), (z[2], t)) = ((1, 2), (20, 4)); [x, y, z, t]",
         [1, 2, [1, 2, 20], 4],
     ],
+    ["a, b, c = [1,2,3]; [a, b, c]", [1, 2, 3]],
+    ["a, b, c = iter([1,2,3]); [a, b, c]", [1, 2, 3]],
+    [
+        "tuples = [(1, 2), (3, 4), (5, 6)]; a, b = zip(*tuples); [a, b]",
+        [(1, 3, 5), (2, 4, 6)],
+    ],
     ["Foo = type('Foo', (), {'x': 100}); Foo.x = 10; Foo.x", 10],
     ["Foo = type('Foo', (), {'x': 100}); Foo.x += 10; Foo.x", 110],
     ["Foo = [type('Foo', (), {'x': 100})]; Foo[0].x = 10; Foo[0].x", 10],
@@ -664,11 +671,19 @@ def test_eval(hass):
     ["xx", "Exception in test line 1 column 0: name 'xx' is not defined"],
     [
         "(x, y) = (1, 2, 4)",
-        "Exception in test line 1 column 16: too many values to unpack (expected 2)",
+        "Exception in test line 1 column 4: too many values to unpack (expected 2)",
+    ],
+    [
+        "(x, y) = iter([1, 2, 4])",
+        "Exception in test line 1 column 4: too many values to unpack (expected 2)",
     ],
     [
         "(x, y, z) = (1, 2)",
-        "Exception in test line 1 column 16: too few values to unpack (expected 3)",
+        "Exception in test line 1 column 4: too few values to unpack (expected 3)",
+    ],
+    [
+        "(x, y, z) = iter([1, 2])",
+        "Exception in test line 1 column 4: too few values to unpack (expected 3)",
     ],
     [
         "(x, y) = 1",