diff --git a/compiler/expr_visitor.py b/compiler/expr_visitor.py index 267abdf7..fb9d2825 100644 --- a/compiler/expr_visitor.py +++ b/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 = { diff --git a/compiler/stmt.py b/compiler/stmt.py index 0b25f32f..967b63d6 100644 --- a/compiler/stmt.py +++ b/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') diff --git a/runtime/generator.go b/runtime/generator.go index 2d2bba66..403cc9fe 100644 --- a/runtime/generator.go +++ b/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} diff --git a/runtime/generator_test.go b/runtime/generator_test.go index 7edbb8be..573caeef 100644 --- a/runtime/generator_test.go +++ b/runtime/generator_test.go @@ -21,7 +21,7 @@ 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 @@ -29,7 +29,7 @@ func TestGeneratorNext(t *testing.T) { 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) + } + } +} diff --git a/testing/generator_test.py b/testing/generator_test.py index 182249c5..37f10b98 100644 --- a/testing/generator_test.py +++ b/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')