blueloveTH 2 years ago
parent
commit
755140084a
5 changed files with 30 additions and 10 deletions
  1. 5 5
      src/ceval.h
  2. 5 5
      src/iter.h
  3. 5 0
      src/pocketpy.h
  4. 3 0
      src/vm.h
  5. 12 0
      tests/28_iter.py

+ 5 - 5
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; 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();
         }

+ 5 - 5
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;
     }
 }
 

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

+ 3 - 0
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<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 - 0
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