Implement generator.throw()

compiler/expr_visitor.py

@@ -379,7 +379,15 @@ def visit_Yield(self, node):
     self.writer.write('return {}, nil'.format(value.expr))
     self.writer.write_label(resume_label)
     result = self.block.alloc_temp()
-    self.writer.write('{} = πSent'.format(result.name))
+    self.writer.write_tmpl(textwrap.dedent("""\
+      if πThrown != nil {
+      \tπE = πThrown
+      \tπThrown = nil
+      \tif πE != nil {
+      \t\tcontinue
+      \t}
+      }
+      $result = πSent"""), result=result.name)
     return result
 
   _BIN_OP_TEMPLATES = {

compiler/stmt.py

@@ -590,9 +590,14 @@ def visit_function_inline(self, node):
         self.writer.write('var πE *πg.BaseException; _ = πE')
         if func_block.is_generator:
           self.writer.write(
-              'return πg.NewGenerator(πF, func(πSent *πg.Object) '
+              'return πg.NewGenerator(πF, '
+              'func(πSent *πg.Object, πThrown *πg.BaseException) '
               '(*πg.Object, *πg.BaseException) {')
           with self.writer.indent_block():
+            self.writer.write(textwrap.dedent("""\
+              if πThrown != nil {
+              \treturn nil, πThrown
+              }"""))
             self.writer.write_block(func_block, visitor.writer.getvalue())
             self.writer.write('return nil, πE')
           self.writer.write('}).ToObject(), nil')

runtime/generator.go

@@ -39,11 +39,11 @@ type Generator struct {
 	mutex sync.Mutex
 	state generatorState
 	frame *Frame
-	fn    func(*Object) (*Object, *BaseException)
+	fn    func(*Object, *BaseException) (*Object, *BaseException)
 }
 
 // NewGenerator returns a new Generator object that runs the given Block b.
-func NewGenerator(f *Frame, fn func(*Object) (*Object, *BaseException)) *Generator {
+func NewGenerator(f *Frame, fn func(*Object, *BaseException) (*Object, *BaseException)) *Generator {
 	f.taken = true // Claim the frame from being returned.
 
 	// The code generator basically gives us the Frame, so we can tare it
@@ -57,13 +57,13 @@ func toGeneratorUnsafe(o *Object) *Generator {
 	return (*Generator)(o.toPointer())
 }
 
-func (g *Generator) resume(f *Frame, sendValue *Object) (*Object, *BaseException) {
+func (g *Generator) resume(f *Frame, sendValue *Object, raisedValue *BaseException) (*Object, *BaseException) {
 	var raised *BaseException
 	g.mutex.Lock()
 	oldState := g.state
 	switch oldState {
 	case generatorStateCreated:
-		if sendValue != None {
+		if sendValue != None && raisedValue == nil {
 			raised = f.RaiseType(TypeErrorType, "can't send non-None value to a just-started generator")
 		} else {
 			g.state = generatorStateRunning
@@ -83,7 +83,7 @@ func (g *Generator) resume(f *Frame, sendValue *Object) (*Object, *BaseException
 		return nil, raised
 	}
 	g.frame.pushFrame(f)
-	result, raised := g.fn(sendValue)
+	result, raised := g.fn(sendValue, raisedValue)
 	g.mutex.Lock()
 	if result == nil && raised == nil {
 		raised = f.Raise(StopIterationType.ToObject(), nil, nil)
@@ -108,18 +108,39 @@ func generatorIter(f *Frame, o *Object) (*Object, *BaseException) {
 }
 
 func generatorNext(f *Frame, o *Object) (*Object, *BaseException) {
-	return toGeneratorUnsafe(o).resume(f, None)
+	return toGeneratorUnsafe(o).resume(f, None, nil)
 }
 
 func generatorSend(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
 	if raised := checkMethodArgs(f, "send", args, GeneratorType, ObjectType); raised != nil {
 		return nil, raised
 	}
-	return toGeneratorUnsafe(args[0]).resume(f, args[1])
+	return toGeneratorUnsafe(args[0]).resume(f, args[1], nil)
+}
+
+func generatorThrow(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
+	argc := len(args)
+	expectedTypes := []*Type{GeneratorType, ObjectType, ObjectType, ObjectType}
+	if argc > 1 && argc < 4 {
+		expectedTypes = expectedTypes[:argc]
+	}
+	if raised := checkMethodArgs(f, "throw", args, expectedTypes...); raised != nil {
+		return nil, raised
+	}
+	var v *Object
+	if argc > 2 {
+		v = args[2]
+	}
+	var tb *Object
+	if argc > 3 {
+		tb = args[3]
+	}
+	return toGeneratorUnsafe(args[0]).resume(f, nil, f.Raise(args[1], v, tb))
 }
 
 func initGeneratorType(dict map[string]*Object) {
 	dict["send"] = newBuiltinFunction("send", generatorSend).ToObject()
+	dict["throw"] = newBuiltinFunction("throw", generatorThrow).ToObject()
 	GeneratorType.flags &= ^(typeFlagBasetype | typeFlagInstantiable)
 	GeneratorType.slots.Iter = &unaryOpSlot{generatorIter}
 	GeneratorType.slots.Next = &unaryOpSlot{generatorNext}

runtime/generator_test.go

@@ -21,15 +21,15 @@ import (
 func TestGeneratorNext(t *testing.T) {
 	f := NewRootFrame()
 	var recursive *Object
-	recursiveFn := func(*Object) (*Object, *BaseException) {
+	recursiveFn := func(*Object, *BaseException) (*Object, *BaseException) {
 		next, raised := GetAttr(f, recursive, NewStr("next"), nil)
 		if raised != nil {
 			return nil, raised
 		}
 		return next.Call(f, nil, nil)
 	}
 	recursive = NewGenerator(f, recursiveFn).ToObject()
-	emptyFn := func(*Object) (*Object, *BaseException) {
+	emptyFn := func(*Object, *BaseException) (*Object, *BaseException) {
 		return nil, nil
 	}
 	exhausted := NewGenerator(NewRootFrame(), emptyFn).ToObject()
@@ -46,7 +46,7 @@ func TestGeneratorNext(t *testing.T) {
 }
 
 func TestGeneratorSend(t *testing.T) {
-	emptyFn := func(*Object) (*Object, *BaseException) {
+	emptyFn := func(*Object, *BaseException) (*Object, *BaseException) {
 		return nil, nil
 	}
 	cases := []invokeTestCase{
@@ -62,7 +62,7 @@ func TestGeneratorSend(t *testing.T) {
 
 func TestGeneratorSimple(t *testing.T) {
 	f := NewRootFrame()
-	fn := func(*Object) (*Object, *BaseException) {
+	fn := func(*Object, *BaseException) (*Object, *BaseException) {
 		switch f.State() {
 		case 0:
 			goto Start
@@ -90,3 +90,27 @@ func TestGeneratorSimple(t *testing.T) {
 		t.Error(err)
 	}
 }
+
+func TestGeneratorThrow(t *testing.T) {
+	emptyFn := func(*Object, *BaseException) (*Object, *BaseException) {
+		return nil, nil
+	}
+	yieldedFn := func(*Object, *BaseException) (*Object, *BaseException) {
+		return NewStr("foo").ToObject(), nil
+	}
+	raisedFn := func(*Object, *BaseException) (*Object, *BaseException) {
+		return nil, NewRootFrame().RaiseType(ValueErrorType, "bar")
+	}
+	cases := []invokeTestCase{
+		invokeTestCase{args: wrapArgs(NewGenerator(NewRootFrame(), emptyFn), TypeErrorType.ToObject()), wantExc: toBaseExceptionUnsafe(mustNotRaise(StopIterationType.Call(NewRootFrame(), nil, nil)))},
+		invokeTestCase{args: wrapArgs(NewGenerator(NewRootFrame(), yieldedFn), TypeErrorType.ToObject()), want: NewStr("foo").ToObject()},
+		invokeTestCase{args: wrapArgs(NewGenerator(NewRootFrame(), raisedFn), TypeErrorType.ToObject()), wantExc: mustCreateException(ValueErrorType, "bar")},
+		invokeTestCase{args: wrapArgs(NewGenerator(NewRootFrame(), emptyFn)), wantExc: mustCreateException(TypeErrorType, "'throw' of 'generator' requires 4 arguments")},
+		invokeTestCase{args: wrapArgs(NewGenerator(NewRootFrame(), emptyFn), "foo", "bar", "baz", "qux"), wantExc: mustCreateException(TypeErrorType, "'throw' of 'generator' requires 4 arguments")},
+	}
+	for _, cas := range cases {
+		if err := runInvokeMethodTestCase(GeneratorType, "throw", &cas); err != "" {
+			t.Error(err)
+		}
+	}
+}

testing/generator_test.py

@@ -74,3 +74,27 @@ def gen6():
 g = gen6()
 assert list(g) == [1]
 assert list(g) == []
+
+
+def gen7():
+  yield
+g = gen7()
+try:
+  g.throw(TypeError, 'foo')
+except TypeError as e:
+  assert "foo" in str(e)
+else:
+  raise AssertionError
+
+
+def gen8():
+  try:
+    yield
+  except ValueError as e:
+    assert "foo" in str(e)
+    yield
+  else:
+    raise AssertionError
+g = gen8()
+g.next()
+g.throw(ValueError, 'foo')