[pypy-commit] pypy py3.3: Implement "yield from" opcode

[amauryfa]

Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: py3.3
Changeset: r72339:02025d9b1f7d
Date: 2014-06-22 19:08 +0200
http://bitbucket.org/pypy/pypy/changeset/02025d9b1f7d/

Log:	Implement "yield from" opcode

diff --git a/pypy/interpreter/astcompiler/codegen.py b/pypy/interpreter/astcompiler/codegen.py
--- a/pypy/interpreter/astcompiler/codegen.py
+++ b/pypy/interpreter/astcompiler/codegen.py
@@ -871,11 +871,12 @@
             self.load_const(self.space.w_None)
         self.emit_op(ops.YIELD_VALUE)
 
-    def visit_YieldFrom(self, yie):
-        # XXX not correctly implemented.
-        self.update_position(yie.lineno)
-        yie.value.walkabout(self)
-        self.emit_op(ops.YIELD_VALUE)
+    def visit_YieldFrom(self, yfr):
+        self.update_position(yfr.lineno)
+        yfr.value.walkabout(self)
+        self.emit_op(ops.GET_ITER)
+        self.load_const(self.space.w_None)
+        self.emit_op(ops.YIELD_FROM)
 
     def visit_Num(self, num):
         self.update_position(num.lineno)
diff --git a/pypy/interpreter/astcompiler/symtable.py b/pypy/interpreter/astcompiler/symtable.py
--- a/pypy/interpreter/astcompiler/symtable.py
+++ b/pypy/interpreter/astcompiler/symtable.py
@@ -431,6 +431,10 @@
         self.scope.note_yield(yie)
         ast.GenericASTVisitor.visit_Yield(self, yie)
 
+    def visit_YieldFrom(self, yfr):
+        self.scope.note_yield(yfr)
+        ast.GenericASTVisitor.visit_YieldFrom(self, yfr)
+
     def visit_Global(self, glob):
         for name in glob.names:
             old_role = self.scope.lookup_role(name)
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
@@ -928,7 +928,6 @@
         expr = self.get_first_expr("yield")
         assert isinstance(expr, ast.Yield)
         assert expr.value is None
-        assert expr.is_from == 0
         expr = self.get_first_expr("yield x")
         assert isinstance(expr.value, ast.Name)
         assign = self.get_first_stmt("x = yield x")
@@ -937,8 +936,8 @@
 
     def test_yield_from(self):
         expr = self.get_first_expr("yield from x")
+        assert isinstance(expr, ast.YieldFrom)
         assert isinstance(expr.value, ast.Name)
-        assert expr.is_from == 1
 
     def test_unaryop(self):
         unary_ops = (
diff --git a/pypy/interpreter/astcompiler/test/test_compiler.py b/pypy/interpreter/astcompiler/test/test_compiler.py
--- a/pypy/interpreter/astcompiler/test/test_compiler.py
+++ b/pypy/interpreter/astcompiler/test/test_compiler.py
@@ -965,6 +965,15 @@
         yield self.st, 'x = list(d for d in [1] or [])', 'x', [1]
         yield self.st, 'y = [d for d in [1] or []]', 'y', [1]
 
+    def test_yield_from(self):
+        test = """if 1:
+        def f():
+            yield from range(3)
+        def g():
+            return list(f())
+        """
+        yield self.st, test, "g()", range(3)
+
 
 class AppTestCompiler:
 
diff --git a/pypy/interpreter/pyopcode.py b/pypy/interpreter/pyopcode.py
--- a/pypy/interpreter/pyopcode.py
+++ b/pypy/interpreter/pyopcode.py
@@ -401,6 +401,8 @@
                 self.WITH_CLEANUP(oparg, next_instr)
             elif opcode == opcodedesc.YIELD_VALUE.index:
                 self.YIELD_VALUE(oparg, next_instr)
+            elif opcode == opcodedesc.YIELD_FROM.index:
+                self.YIELD_FROM(oparg, next_instr)
             else:
                 self.MISSING_OPCODE(oparg, next_instr)
 
@@ -1000,6 +1002,34 @@
     def YIELD_VALUE(self, oparg, next_instr):
         raise Yield
 
+    def YIELD_FROM(self, oparg, next_instr):
+        space = self.space
+        w_value = self.popvalue()
+        w_gen = self.peekvalue()
+        try:
+            if space.is_none(w_value):
+                w_retval = space.next(w_gen)
+            else:
+                w_retval = space.call_method(w_gen, "send", w_value)
+        except OperationError as e:
+            if not e.match(self.space, self.space.w_StopIteration):
+                raise
+            self.popvalue()  # Remove iter from stack
+            try:
+                w_value = space.getattr(e.get_w_value(space), space.wrap("value"))
+            except OperationError as e:
+                if not e.match(self.space, self.space.w_AttributeError):
+                    raise
+                w_value = space.w_None
+            self.pushvalue(w_value)
+            return next_instr
+        else:
+            # iter remains on stack, w_retval is value to be yielded.
+            self.pushvalue(w_retval)
+            # and repeat...
+            self.last_instr = self.last_instr - 1
+            raise Yield
+
     def jump_absolute(self, jumpto, ec):
         # this function is overridden by pypy.module.pypyjit.interp_jit
         check_nonneg(jumpto)