Merge pull request #592 from abonie/yield_from

Support for `yield from` statement

batavia/VirtualMachine.js

@@ -8,6 +8,7 @@ var Frame = require('./core').Frame
 var constants = require('./core').constants
 var exceptions = require('./core').exceptions
 var native = require('./core').native
+var callables = require('./core').callables
 var dis = require('./modules/dis')
 var marshal = require('./modules/marshal')
 var sys = require('./modules/sys')
@@ -1847,30 +1848,29 @@ VirtualMachine.prototype.byte_YIELD_VALUE = function() {
     return 'yield'
 }
 
-// VirtualMachine.prototype.byte_YIELD_FROM = function {
-//         u = this.pop()
-//         x = this.top()
+VirtualMachine.prototype.byte_YIELD_FROM = function() {
+    var v = this.pop()
+    var receiver = this.top()
 
-//         try:
-//             if not isinstance(x, Generator) or u is null:
-//                 // Call next on iterators.
-//                 retval = next(x)
-//             else:
-//                 retval = x.send(u)
-//             this.return_value = retval
-//         except StopIteration as e:
-//             this.pop()
-//             this.push(e.value)
-//         else:
-//             // YIELD_FROM decrements f_lasti, so that it will be called
-//             // repeatedly until a StopIteration is raised.
-//             this.jump(this.frame.f_lasti - 1)
-//             // Returning "yield" prevents the block stack cleanup code
-//             // from executing, suspending the frame in its current state.
-//             return "yield"
-
-//     #// Importing
-// }
+    try {
+        if (types.isinstance(v, types.NoneType) ||
+                !types.isinstance(receiver, types.Generator)) {
+            this.return_value = callables.call_method(receiver, '__next__', [])
+        } else {
+            this.return_value = receiver.send(v)
+        }
+    } catch (e) {
+        if (e instanceof exceptions.StopIteration.$pyclass) {
+            this.pop()
+            this.push(e.value)
+            return
+        } else {
+            throw e
+        }
+    }
+    this.jump(this.frame.f_lasti - 1)
+    return 'yield'
+}
 
 VirtualMachine.prototype.byte_IMPORT_NAME = function(name) {
     var items = this.popn(2)

batavia/types/Generator.js

@@ -2,6 +2,7 @@ var PyObject = require('../core').Object
 var create_pyclass = require('../core').create_pyclass
 var exceptions = require('../core').exceptions
 var callables = require('../core').callables
+var dis = require('../modules/dis')
 
 /*************************************************************************
  * A Python generator type.
@@ -62,6 +63,21 @@ Generator.prototype.send = function(value) {
 }
 
 Generator.prototype['throw'] = function(ExcType, value, traceback) {
+    var yf_gen = yf_subgenerator(this)
+    if (yf_gen !== null) {
+        if (ExcType instanceof exceptions.GeneratorExit.$pyclass ||
+                value instanceof exceptions.GeneratorExit.$pyclass) {
+            callables.call_method(yf_gen, 'close', [])
+        } else {
+            try {
+                return callables.call_method(yf_gen, 'throw', [ExcType, value, traceback])
+            } catch (e) {
+                if (!(e instanceof exceptions.AttributeError.$pyclass)) {
+                    throw e
+                }
+            }
+        }
+    }
     if (ExcType instanceof exceptions.BaseException.$pyclass) {
         value = ExcType
         ExcType = ExcType.__class__
@@ -98,6 +114,18 @@ Generator.prototype['close'] = function() {
     }
 }
 
+const YIELD_FROM = dis.opmap['YIELD_FROM']
+
+// returns subgenerator gen is yielding from or null if there is none
+function yf_subgenerator(gen) {
+    var f = gen.gi_frame
+    var opcode = f.f_code.co_code.valueOf()[f.f_lasti]
+    if (opcode === YIELD_FROM) {
+        return f.stack[1]
+    }
+    return null
+}
+
 /**************************************************
  * Module exports
  **************************************************/

tests/structures/test_generator.py

@@ -193,3 +193,200 @@ def G():
             except BaseException as e:
                 print(type(e), e)
         """)
+
+    def test_nested(self):
+        self.assertCodeExecution("""
+            def G():
+                while True:
+                    try:
+                        v = (yield)
+                    except KeyError:
+                        print('ERROR')
+                    else:
+                        print('>>', v)
+
+            def F(g):
+                g.send(None)
+                while True:
+                    try:
+                        try:
+                            v = (yield)
+                        except Exception as e:
+                            g.throw(e)
+                        else:
+                            g.send(v)
+                    except StopIteration:
+                        pass
+
+            g = G()
+            g2 = F(g)
+            g2.send(None)
+            g2.send(1)
+            g2.send(2)
+            g2.throw(KeyError)
+            g2.send(3)
+        """)
+
+
+class YieldFromTests(TranspileTestCase):
+    def test_simple_generator(self):
+        self.assertCodeExecution("""
+            def G(x):
+                yield from x
+
+            def F():
+                yield 1
+                yield 2
+
+            g = G(F())
+            print(list(g))
+        """)
+
+    def test_yield_before_and_after(self):
+        self.assertCodeExecution("""
+            def G(x):
+                yield "START"
+                yield from x
+                yield "STOP"
+
+            def F():
+                yield 1
+                yield 2
+
+            g = G(F())
+            print(list(g))
+        """)
+
+    def test_iterator(self):
+        self.assertCodeExecution("""
+            def G(x):
+                yield "START"
+                yield from x
+                yield "STOP"
+
+            g = G(range(5))
+            print(list(g))
+        """)
+
+    def test_throw(self):
+        self.assertCodeExecution("""
+            def G():
+                yield 1
+                yield 2
+                yield 3
+
+            def F(g):
+                yield from g
+
+            f = F(G())
+            print(next(f))
+            try:
+                f.throw(KeyError)
+            except Exception as e:
+                print(type(e), e)
+            try:
+                print(next(f))
+            except Exception as e:
+                print(type(e), e)
+        """)
+
+    def test_catch_exception(self):
+        self.assertCodeExecution("""
+            def G():
+                while True:
+                    try:
+                        v = (yield)
+                    except KeyError:
+                        print('ERROR')
+                    else:
+                        print('>>', v)
+
+            def F(g):
+                yield from g
+
+            g = F(G())
+            g.send(None)
+            g.send(1)
+            g.send(2)
+            g.throw(KeyError)
+            g.send(3)
+        """)
+
+    def test_throw_iterator(self):
+        self.assertCodeExecution("""
+            def G(g):
+                yield from g
+
+            g = G(range(5))
+            next(g)
+            try:
+                g.throw(KeyError)
+            except Exception as e:
+                print(type(e), e)
+        """)
+
+    def test_pyclass_subgenerator(self):
+        self.assertCodeExecution("""
+            class Gen:
+                def __init__(self):
+                    self.stopped = False
+                    self.i = -1
+                def __next__(self):
+                    if not self.stopped:
+                        self.i += 1
+                        return self.i
+                    raise StopIteration
+                def __iter__(self):
+                    return self
+                def close(self):
+                    self.stopped = True
+
+            def G(g):
+                yield 'START'
+                yield from g
+                yield 'STOP'
+
+            g = Gen()
+            gg = G(g)
+            print(next(gg))
+            print(next(gg))
+            print(gg.close())
+            try:
+                print(next(gg))
+            except StopIteration as e:
+                print(type(e), e)
+        """)
+
+    def test_pyclass_custom_throw(self):
+        self.assertCodeExecution("""
+            class Gen:
+                def __init__(self):
+                    self.stopped = False
+                    self.i = -1
+                def __next__(self):
+                    if not self.stopped:
+                        self.i += 1
+                        return self.i
+                    raise StopIteration
+                def __iter__(self):
+                    return self
+                def close(self):
+                    self.stopped = True
+                def throw(self, a, b, c):
+                    print('THROW')
+
+            def G(g):
+                yield 'START'
+                yield from g
+                yield 'STOP'
+
+            g = Gen()
+            gg = G(g)
+            print(next(gg))
+            print(next(gg))
+            print(gg.close())
+            try:
+                print(next(gg))
+            except StopIteration as e:
+                print(type(e), e)
+        """)