[pypy-commit] pypy py3k: (antocuni, romain) Add support for the new unpacking at the ast level

[rguillebert]

Author: Romain Guillebert <romain.py at gmail.com>
Branch: py3k
Changeset: r51410:fbf25d09b601
Date: 2012-01-17 18:51 +0100
http://bitbucket.org/pypy/pypy/changeset/fbf25d09b601/

Log:	(antocuni,romain) Add support for the new unpacking at the ast level
	fixed list comprehension

diff --git a/pypy/interpreter/astcompiler/ast.py b/pypy/interpreter/astcompiler/ast.py
--- a/pypy/interpreter/astcompiler/ast.py
+++ b/pypy/interpreter/astcompiler/ast.py
@@ -1507,6 +1507,11 @@
             for node in self.comparators:
                 node.sync_app_attrs(space)
 
+class Starred(expr):
+    def __init__(self, value, ctx, lineno, col_offset):
+        self.value = value
+        self.ctx = ctx
+        expr.__init__(self, lineno, col_offset)
 
 class Call(expr):
 
diff --git a/pypy/interpreter/astcompiler/astbuilder.py b/pypy/interpreter/astcompiler/astbuilder.py
--- a/pypy/interpreter/astcompiler/astbuilder.py
+++ b/pypy/interpreter/astcompiler/astbuilder.py
@@ -680,7 +680,7 @@
                 self.set_context(target_expr, ast.Store)
                 targets.append(target_expr)
             value_child = stmt.children[-1]
-            if value_child.type == syms.testlist or value_child.type == syms.testlist_star_expr:
+            if value_child.type == syms.testlist_star_expr:
                 value_expr = self.handle_testlist(value_child)
             else:
                 value_expr = self.handle_expr(value_child)
@@ -740,6 +740,8 @@
                     operands.append(self.handle_expr(expr_node.children[i + 1]))
                 return ast.Compare(expr, operators, operands, expr_node.lineno,
                                    expr_node.column)
+            elif expr_node_type == syms.star_expr:
+                return self.handle_star_expr(expr_node)
             elif expr_node_type == syms.expr or \
                     expr_node_type == syms.xor_expr or \
                     expr_node_type == syms.and_expr or \
@@ -766,6 +768,10 @@
             else:
                 raise AssertionError("unknown expr")
 
+    def handle_star_expr(self, star_expr_node):
+        expr = self.handle_expr(star_expr_node.children[1])
+        return ast.Starred(expr, ast.Load, star_expr_node.lineno, star_expr_node.column)
+
     def handle_lambdef(self, lambdef_node):
         expr = self.handle_expr(lambdef_node.children[-1])
         if len(lambdef_node.children) == 3:
@@ -1229,8 +1235,8 @@
         elt = self.handle_expr(listcomp_node.children[0])
         comps = self.comprehension_helper(listcomp_node.children[1],
                                           "handle_testlist",
-                                          syms.list_for, syms.list_if,
-                                          syms.list_iter,
+                                          syms.comp_for, syms.comp_if,
+                                          syms.comp_iter,
                                           comp_fix_unamed_tuple_location=True)
         return ast.ListComp(elt, comps, listcomp_node.lineno,
                             listcomp_node.column)
diff --git a/pypy/interpreter/astcompiler/asthelpers.py b/pypy/interpreter/astcompiler/asthelpers.py
--- a/pypy/interpreter/astcompiler/asthelpers.py
+++ b/pypy/interpreter/astcompiler/asthelpers.py
@@ -133,6 +133,13 @@
 
     _description = "comparison"
 
+class __extend__(ast.Starred):
+
+    _description = "starred expression"
+
+    def set_context(self, ctx):
+        self.ctx = ctx
+        self.value.set_context(ctx)
 
 class __extend__(ast.IfExp):
 
diff --git a/pypy/interpreter/astcompiler/test/test_astbuilder.py b/pypy/interpreter/astcompiler/test/test_astbuilder.py
--- a/pypy/interpreter/astcompiler/test/test_astbuilder.py
+++ b/pypy/interpreter/astcompiler/test/test_astbuilder.py
@@ -600,6 +600,17 @@
         assert isinstance(tup.elts[0], ast.Name)
         assert tup.elts[0].ctx == ast.Store
 
+    def test_assign_starred(self):
+        assign = self.get_first_stmt("*a, b = x")
+        assert isinstance(assign, ast.Assign)
+        assert len(assign.targets) == 1
+        names = assign.targets[0]
+        assert len(names.elts) == 2
+        assert isinstance(names.elts[0], ast.Starred)
+        assert isinstance(names.elts[1], ast.Name)
+        assert isinstance(names.elts[0].value, ast.Name)
+        assert names.elts[0].value.id == "a"
+
     def test_name(self):
         name = self.get_first_expr("hi")
         assert isinstance(name, ast.Name)
diff --git a/pypy/interpreter/pyparser/data/Grammar3.2 b/pypy/interpreter/pyparser/data/Grammar3.2
--- a/pypy/interpreter/pyparser/data/Grammar3.2
+++ b/pypy/interpreter/pyparser/data/Grammar3.2
@@ -127,11 +127,6 @@
 # The reason that keywords are test nodes instead of NAME is that using NAME
 # results in an ambiguity. ast.c makes sure it's a NAME.
 argument: test [comp_for] | test '=' test
-
-list_iter: list_for | list_if
-list_for: 'for' exprlist 'in' testlist_safe [list_iter]
-list_if: 'if' old_test [list_iter]
-
 comp_iter: comp_for | comp_if
 comp_for: 'for' exprlist 'in' or_test [comp_iter]
 comp_if: 'if' old_test [comp_iter]