diff --git a/src/iter.cpp b/src/iter.cpp index e62108f1..3f2d6be6 100644 --- a/src/iter.cpp +++ b/src/iter.cpp @@ -7,11 +7,15 @@ namespace pkpy{ vm->bind_notimplemented_constructor(type); vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* obj){ return obj; }); vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* obj){ - RangeIter& self = _CAST(RangeIter&, obj); - bool has_next = self.r.step > 0 ? self.current < self.r.stop : self.current > self.r.stop; - if(!has_next) return vm->StopIteration; + RangeIter& self = PK_OBJ_GET(RangeIter, obj); + if(self.r.step > 0){ + if(self.current >= self.r.stop) return vm->StopIteration; + }else{ + if(self.current <= self.r.stop) return vm->StopIteration; + } + PyObject* ret = VAR(self.current); self.current += self.r.step; - return VAR(self.current - self.r.step); + return ret; }); } diff --git a/src/pocketpy.cpp b/src/pocketpy.cpp index 53e61949..6c35555a 100644 --- a/src/pocketpy.cpp +++ b/src/pocketpy.cpp @@ -352,6 +352,7 @@ void init_builtins(VM* _vm) { case 3: r.start = CAST(i64, args[0]); r.stop = CAST(i64, args[1]); r.step = CAST(i64, args[2]); break; default: vm->TypeError("expected 1-3 arguments, got " + std::to_string(args.size())); } + if(r.step == 0) vm->ValueError("range() arg 3 must not be zero"); return VAR(r); });