diff --git a/src/ceval.h b/src/ceval.h index 81ab9c57..23ae3e46 100644 --- a/src/ceval.h +++ b/src/ceval.h @@ -426,7 +426,7 @@ __NEXT_STEP:; BaseIter* it = _PyIter_AS_C(TOP()); #endif PyObject* obj = it->next(); - if(obj != nullptr){ + if(obj != StopIteration){ PUSH(obj); }else{ int target = co_blocks[byte.block].end; @@ -474,7 +474,7 @@ __NEXT_STEP:; BaseIter* iter = PyIter_AS_C(obj); for(int i=0; inext(); - if(item == nullptr) ValueError("not enough values to unpack"); + if(item == StopIteration) ValueError("not enough values to unpack"); PUSH(item); } // handle extra items @@ -482,12 +482,12 @@ __NEXT_STEP:; List extras; while(true){ PyObject* item = iter->next(); - if(item == nullptr) break; + if(item == StopIteration) break; extras.push_back(item); } PUSH(VAR(extras)); }else{ - if(iter->next() != nullptr) ValueError("too many values to unpack"); + if(iter->next() != StopIteration) ValueError("too many values to unpack"); } } DISPATCH(); TARGET(UNPACK_UNLIMITED) { @@ -495,7 +495,7 @@ __NEXT_STEP:; PyObject* obj = asIter(POPX()); BaseIter* iter = PyIter_AS_C(obj); obj = iter->next(); - while(obj != nullptr){ + while(obj != StopIteration){ PUSH(obj); obj = iter->next(); } diff --git a/src/iter.h b/src/iter.h index f06047ec..b392c4aa 100644 --- a/src/iter.h +++ b/src/iter.h @@ -19,7 +19,7 @@ public: } PyObject* next(){ - if(!_has_next()) return nullptr; + if(!_has_next()) return vm->StopIteration; current += r.step; return VAR(current-r.step); } @@ -37,7 +37,7 @@ public: } PyObject* next() override{ - if(index >= array->size()) return nullptr; + if(index >= array->size()) return vm->StopIteration; return array->operator[](index++); } @@ -56,7 +56,7 @@ public: // TODO: optimize this to use iterator // operator[] is O(n) complexity Str* str = &OBJ_GET(Str, ref); - if(index == str->u8_length()) return nullptr; + if(index == str->u8_length()) return vm->StopIteration; return VAR(str->u8_getitem(index++)); } @@ -66,7 +66,7 @@ public: }; inline PyObject* Generator::next(){ - if(state == 2) return nullptr; + if(state == 2) return vm->StopIteration; // reset frame._sp_base frame._sp_base = frame._s->_sp; frame._locals.a = frame._s->_sp; @@ -85,7 +85,7 @@ inline PyObject* Generator::next(){ return ret; }else{ state = 2; - return nullptr; + return vm->StopIteration; } } diff --git a/src/pocketpy.h b/src/pocketpy.h index 4886be57..4311600d 100644 --- a/src/pocketpy.h +++ b/src/pocketpy.h @@ -163,6 +163,11 @@ inline void init_builtins(VM* _vm) { return vm->asIter(args[0]); }); + _vm->bind_builtin_func<1>("next", [](VM* vm, ArgsView args) { + BaseIter* iter = vm->PyIter_AS_C(args[0]); + return iter->next(); + }); + _vm->bind_builtin_func<1>("dir", [](VM* vm, ArgsView args) { std::set names; if(args[0]->is_attr_valid()){ diff --git a/src/vm.h b/src/vm.h index 3d542714..37899ad4 100644 --- a/src/vm.h +++ b/src/vm.h @@ -87,6 +87,7 @@ public: PyObject* False; PyObject* Ellipsis; PyObject* builtins; // builtins module + PyObject* StopIteration; PyObject* _main; // __main__ module std::stringstream _stdout_buffer; @@ -745,6 +746,7 @@ inline void VM::init_builtin_types(){ this->Ellipsis = heap._new(_new_type_object("ellipsis"), {}); this->True = heap._new(tp_bool, {}); this->False = heap._new(tp_bool, {}); + this->StopIteration = heap._new(_new_type_object("StopIterationType"), {}); this->builtins = new_module("builtins"); this->_main = new_module("__main__"); @@ -759,6 +761,7 @@ inline void VM::init_builtin_types(){ builtins->attr().set("list", _t(tp_list)); builtins->attr().set("tuple", _t(tp_tuple)); builtins->attr().set("range", _t(tp_range)); + builtins->attr().set("StopIteration", StopIteration); post_init(); for(int i=0; i<_all_types.size(); i++){ diff --git a/tests/28_iter.py b/tests/28_iter.py new file mode 100644 index 00000000..49786af6 --- /dev/null +++ b/tests/28_iter.py @@ -0,0 +1,12 @@ +a = [1, 2, 3] +a = iter(a) + +total = 0 + +while True: + obj = next(a) + if obj is StopIteration: + break + total += obj + +assert total == 6