This commit is contained in:
blueloveTH 2023-04-19 13:25:34 +08:00
parent c1b2e3e52b
commit 755140084a
5 changed files with 30 additions and 10 deletions

View File

@ -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; i<byte.arg; i++){
PyObject* item = iter->next();
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();
}

View File

@ -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;
}
}

View File

@ -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<StrName> names;
if(args[0]->is_attr_valid()){

View File

@ -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<Dummy>(_new_type_object("ellipsis"), {});
this->True = heap._new<Dummy>(tp_bool, {});
this->False = heap._new<Dummy>(tp_bool, {});
this->StopIteration = heap._new<Dummy>(_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++){

12
tests/28_iter.py Normal file
View File

@ -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