github.com/google/grumpy@v0.0.0-20171122020858-3ec87959189c/runtime/range.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 "fmt" 19 "reflect" 20 "sync" 21 ) 22 23 var ( 24 // enumerateType is the object representing the Python 'enumerate' type. 25 enumerateType = newBasisType("enumerate", reflect.TypeOf(enumerate{}), toEnumerateUnsafe, ObjectType) 26 // rangeIteratorType is the object representing the Python 'rangeiterator' type. 27 rangeIteratorType = newBasisType("rangeiterator", reflect.TypeOf(rangeIterator{}), toRangeIteratorUnsafe, ObjectType) 28 // xrangeType is the object representing the Python 'xrange' type. 29 xrangeType = newBasisType("xrange", reflect.TypeOf(xrange{}), toXRangeUnsafe, ObjectType) 30 ) 31 32 type enumerate struct { 33 Object 34 mutex sync.Mutex 35 index int 36 iter *Object 37 } 38 39 func toEnumerateUnsafe(o *Object) *enumerate { 40 return (*enumerate)(o.toPointer()) 41 } 42 43 func enumerateIter(f *Frame, o *Object) (*Object, *BaseException) { 44 return o, nil 45 } 46 47 func enumerateNew(f *Frame, t *Type, args Args, _ KWArgs) (*Object, *BaseException) { 48 expectedTypes := []*Type{ObjectType, ObjectType} 49 argc := len(args) 50 if argc == 1 { 51 expectedTypes = expectedTypes[:1] 52 } 53 if raised := checkFunctionArgs(f, "__new__", args, expectedTypes...); raised != nil { 54 return nil, raised 55 } 56 index := 0 57 if argc > 1 { 58 if args[1].typ.slots.Index == nil { 59 format := "%s object cannot be interpreted as an index" 60 return nil, f.RaiseType(TypeErrorType, fmt.Sprintf(format, args[1].typ.Name())) 61 } 62 // TODO: support long? 63 i, raised := IndexInt(f, args[1]) 64 if raised != nil { 65 return nil, raised 66 } 67 if i > 0 { 68 index = i 69 } 70 } 71 iter, raised := Iter(f, args[0]) 72 if raised != nil { 73 return nil, raised 74 } 75 for i := 0; i < index; i++ { 76 _, raised := Next(f, iter) 77 if raised != nil { 78 if !raised.isInstance(StopIterationType) { 79 return nil, raised 80 } 81 index = -1 82 f.RestoreExc(nil, nil) 83 break 84 } 85 } 86 var d *Dict 87 if t != enumerateType { 88 d = NewDict() 89 } 90 e := &enumerate{Object: Object{typ: t, dict: d}, index: index, iter: iter} 91 return &e.Object, nil 92 } 93 94 func enumerateNext(f *Frame, o *Object) (ret *Object, raised *BaseException) { 95 e := toEnumerateUnsafe(o) 96 e.mutex.Lock() 97 var item *Object 98 if e.index != -1 { 99 item, raised = Next(f, e.iter) 100 } 101 if raised == nil { 102 if item == nil { 103 raised = f.Raise(StopIterationType.ToObject(), nil, nil) 104 e.index = -1 105 } else { 106 ret = NewTuple2(NewInt(e.index).ToObject(), item).ToObject() 107 e.index++ 108 } 109 } 110 e.mutex.Unlock() 111 return ret, raised 112 } 113 114 func initEnumerateType(map[string]*Object) { 115 enumerateType.slots.Iter = &unaryOpSlot{enumerateIter} 116 enumerateType.slots.Next = &unaryOpSlot{enumerateNext} 117 enumerateType.slots.New = &newSlot{enumerateNew} 118 } 119 120 // TODO: Synchronize access to this structure. 121 type rangeIterator struct { 122 Object 123 i int 124 stop int 125 step int 126 } 127 128 func toRangeIteratorUnsafe(o *Object) *rangeIterator { 129 return (*rangeIterator)(o.toPointer()) 130 } 131 132 func rangeIteratorIter(f *Frame, o *Object) (*Object, *BaseException) { 133 return o, nil 134 } 135 136 func rangeIteratorNext(f *Frame, o *Object) (*Object, *BaseException) { 137 iter := toRangeIteratorUnsafe(o) 138 if iter.i == iter.stop { 139 return nil, f.Raise(StopIterationType.ToObject(), nil, nil) 140 } 141 ret := NewInt(iter.i) 142 iter.i += iter.step 143 return ret.ToObject(), nil 144 } 145 146 func initRangeIteratorType(map[string]*Object) { 147 rangeIteratorType.flags &^= typeFlagInstantiable | typeFlagBasetype 148 rangeIteratorType.slots.Iter = &unaryOpSlot{rangeIteratorIter} 149 rangeIteratorType.slots.Next = &unaryOpSlot{rangeIteratorNext} 150 } 151 152 type xrange struct { 153 Object 154 start int 155 stop int 156 step int 157 } 158 159 func toXRangeUnsafe(o *Object) *xrange { 160 return (*xrange)(o.toPointer()) 161 } 162 163 func xrangeGetItem(f *Frame, o, key *Object) (*Object, *BaseException) { 164 if key.typ.slots.Index == nil { 165 format := "sequence index must be integer, not '%s'" 166 return nil, f.RaiseType(TypeErrorType, fmt.Sprintf(format, key.typ.Name())) 167 } 168 i, raised := IndexInt(f, key) 169 if raised != nil { 170 return nil, raised 171 } 172 r := toXRangeUnsafe(o) 173 i, raised = seqCheckedIndex(f, (r.stop-r.start)/r.step, i) 174 if raised != nil { 175 return nil, raised 176 } 177 return NewInt(r.start + i*r.step).ToObject(), nil 178 } 179 180 func xrangeIter(f *Frame, o *Object) (*Object, *BaseException) { 181 r := toXRangeUnsafe(o) 182 return &(&rangeIterator{Object{typ: rangeIteratorType}, r.start, r.stop, r.step}).Object, nil 183 } 184 185 func xrangeLen(f *Frame, o *Object) (*Object, *BaseException) { 186 r := toXRangeUnsafe(o) 187 return NewInt((r.stop - r.start) / r.step).ToObject(), nil 188 } 189 190 func xrangeNew(f *Frame, _ *Type, args Args, _ KWArgs) (*Object, *BaseException) { 191 expectedTypes := []*Type{IntType, IntType, IntType} 192 argc := len(args) 193 if argc > 0 && argc < 3 { 194 expectedTypes = expectedTypes[:argc] 195 } 196 if raised := checkMethodArgs(f, "__new__", args, expectedTypes...); raised != nil { 197 return nil, raised 198 } 199 start, stop, step := 0, 0, 1 200 if argc == 1 { 201 stop = toIntUnsafe(args[0]).Value() 202 } else { 203 start = toIntUnsafe(args[0]).Value() 204 stop = toIntUnsafe(args[1]).Value() 205 if argc > 2 { 206 step = toIntUnsafe(args[2]).Value() 207 } 208 } 209 stop, _, result := seqRange(start, stop, step) 210 switch result { 211 case seqRangeZeroStep: 212 return nil, f.RaiseType(ValueErrorType, "xrange() arg 3 must not be zero") 213 case seqRangeOverflow: 214 return nil, f.RaiseType(OverflowErrorType, errResultTooLarge) 215 } 216 r := &xrange{Object: Object{typ: xrangeType}, start: start, stop: stop, step: step} 217 return &r.Object, nil 218 } 219 220 func xrangeRepr(_ *Frame, o *Object) (*Object, *BaseException) { 221 r := toXRangeUnsafe(o) 222 s := "" 223 if r.step != 1 { 224 s = fmt.Sprintf("xrange(%d, %d, %d)", r.start, r.stop, r.step) 225 } else if r.start != 0 { 226 s = fmt.Sprintf("xrange(%d, %d)", r.start, r.stop) 227 } else { 228 s = fmt.Sprintf("xrange(%d)", r.stop) 229 } 230 return NewStr(s).ToObject(), nil 231 } 232 233 func initXRangeType(map[string]*Object) { 234 xrangeType.flags &^= typeFlagBasetype 235 xrangeType.slots.GetItem = &binaryOpSlot{xrangeGetItem} 236 xrangeType.slots.Iter = &unaryOpSlot{xrangeIter} 237 xrangeType.slots.Len = &unaryOpSlot{xrangeLen} 238 xrangeType.slots.New = &newSlot{xrangeNew} 239 xrangeType.slots.Repr = &unaryOpSlot{xrangeRepr} 240 }