github.com/google/grumpy@v0.0.0-20171122020858-3ec87959189c/runtime/generator.go (about) 1 // Copyright 2016 Google Inc. All Rights Reserved. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package grumpy 16 17 import ( 18 "reflect" 19 "sync" 20 ) 21 22 var ( 23 // GeneratorType is the object representing the Python 'generator' type. 24 GeneratorType = newBasisType("generator", reflect.TypeOf(Generator{}), toGeneratorUnsafe, ObjectType) 25 ) 26 27 type generatorState int 28 29 const ( 30 generatorStateCreated generatorState = iota 31 generatorStateReady 32 generatorStateRunning 33 generatorStateDone 34 ) 35 36 // Generator represents Python 'generator' objects. 37 type Generator struct { 38 Object 39 mutex sync.Mutex 40 state generatorState 41 frame *Frame 42 fn func(*Object) (*Object, *BaseException) 43 } 44 45 // NewGenerator returns a new Generator object that runs the given Block b. 46 func NewGenerator(f *Frame, fn func(*Object) (*Object, *BaseException)) *Generator { 47 f.taken = true // Claim the frame from being returned. 48 49 // The code generator basically gives us the Frame, so we can tare it 50 // off and prevent a parasitic `taken` from creeping up the frames. 51 f.back = nil 52 53 return &Generator{Object: Object{typ: GeneratorType}, frame: f, fn: fn} 54 } 55 56 func toGeneratorUnsafe(o *Object) *Generator { 57 return (*Generator)(o.toPointer()) 58 } 59 60 func (g *Generator) resume(f *Frame, sendValue *Object) (*Object, *BaseException) { 61 var raised *BaseException 62 g.mutex.Lock() 63 oldState := g.state 64 switch oldState { 65 case generatorStateCreated: 66 if sendValue != None { 67 raised = f.RaiseType(TypeErrorType, "can't send non-None value to a just-started generator") 68 } else { 69 g.state = generatorStateRunning 70 } 71 case generatorStateReady: 72 g.state = generatorStateRunning 73 case generatorStateRunning: 74 raised = f.RaiseType(ValueErrorType, "generator already executing") 75 case generatorStateDone: 76 raised = f.Raise(StopIterationType.ToObject(), nil, nil) 77 } 78 g.mutex.Unlock() 79 // Concurrent attempts to transition to running state will raise here 80 // so it's guaranteed that only one thread will proceed to execute the 81 // block below. 82 if raised != nil { 83 return nil, raised 84 } 85 g.frame.pushFrame(f) 86 result, raised := g.fn(sendValue) 87 g.mutex.Lock() 88 if result == nil && raised == nil { 89 raised = f.Raise(StopIterationType.ToObject(), nil, nil) 90 } 91 if raised == nil { 92 g.frame.PopCheckpoint() 93 g.state = generatorStateReady 94 } else { 95 g.state = generatorStateDone 96 } 97 g.mutex.Unlock() 98 return result, raised 99 } 100 101 // ToObject upcasts g to an Object. 102 func (g *Generator) ToObject() *Object { 103 return &g.Object 104 } 105 106 func generatorIter(f *Frame, o *Object) (*Object, *BaseException) { 107 return o, nil 108 } 109 110 func generatorNext(f *Frame, o *Object) (*Object, *BaseException) { 111 return toGeneratorUnsafe(o).resume(f, None) 112 } 113 114 func generatorSend(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) { 115 if raised := checkMethodArgs(f, "send", args, GeneratorType, ObjectType); raised != nil { 116 return nil, raised 117 } 118 return toGeneratorUnsafe(args[0]).resume(f, args[1]) 119 } 120 121 func initGeneratorType(dict map[string]*Object) { 122 dict["send"] = newBuiltinFunction("send", generatorSend).ToObject() 123 GeneratorType.flags &= ^(typeFlagBasetype | typeFlagInstantiable) 124 GeneratorType.slots.Iter = &unaryOpSlot{generatorIter} 125 GeneratorType.slots.Next = &unaryOpSlot{generatorNext} 126 }