Drop iter_type_init, explicitly define __iter__ for iterators

Author: adrian17

vm/src/obj/objenumerate.rs

@@ -57,13 +57,17 @@ impl PyEnumerateRef {
 
         Ok(result)
     }
+
+    fn iter(self, _vm: &VirtualMachine) -> Self {
+        self
+    }
 }
 
 pub fn init(context: &PyContext) {
     let enumerate_type = &context.enumerate_type;
-    objiter::iter_type_init(context, enumerate_type);
     extend_class!(context, enumerate_type, {
         "__new__" => context.new_rustfunc(enumerate_new),
-        "__next__" => context.new_rustfunc(PyEnumerateRef::next)
+        "__next__" => context.new_rustfunc(PyEnumerateRef::next),
+        "__iter__" => context.new_rustfunc(PyEnumerateRef::iter),
     });
 }

vm/src/obj/objfilter.rs

@@ -52,13 +52,15 @@ impl PyFilterRef {
             }
         }
     }
+
+    fn iter(self, _vm: &VirtualMachine) -> Self {
+        self
+    }
 }
 
 pub fn init(context: &PyContext) {
     let filter_type = &context.filter_type;
 
-    objiter::iter_type_init(context, filter_type);
-
     let filter_doc =
         "filter(function or None, iterable) --> filter object\n\n\
          Return an iterator yielding those items of iterable for which function(item)\n\
@@ -67,6 +69,7 @@ pub fn init(context: &PyContext) {
     extend_class!(context, filter_type, {
         "__new__" => context.new_rustfunc(filter_new),
         "__doc__" => context.new_str(filter_doc.to_string()),
-        "__next__" => context.new_rustfunc(PyFilterRef::next)
+        "__next__" => context.new_rustfunc(PyFilterRef::next),
+        "__iter__" => context.new_rustfunc(PyFilterRef::iter),
     });
 }

vm/src/obj/objiter.rs

@@ -2,16 +2,14 @@
  * Various types to support iteration.
  */
 
-use crate::function::PyFuncArgs;
-use crate::pyobject::{PyContext, PyIteratorValue, PyObjectRef, PyResult, TypeProtocol};
+use crate::pyobject::{PyContext, PyIteratorValue, PyObjectRef, PyRef, PyResult};
 use crate::vm::VirtualMachine;
 
 use super::objbytearray::PyByteArray;
 use super::objbytes::PyBytes;
 use super::objrange::PyRange;
 use super::objsequence;
 use super::objtype;
-use crate::obj::objtype::PyClassRef;
 
 /*
  * This helper function is called at multiple places. First, it is called
@@ -73,41 +71,12 @@ pub fn new_stop_iteration(vm: &VirtualMachine) -> PyObjectRef {
     vm.new_exception(stop_iteration_type, "End of iterator".to_string())
 }
 
-/// Common setup for iter types, adds __iter__ method
-pub fn iter_type_init(context: &PyContext, iter_type: &PyClassRef) {
-    let iter_func = {
-        let cloned_iter_type = iter_type.clone();
-        move |vm: &VirtualMachine, args: PyFuncArgs| {
-            arg_check!(
-                vm,
-                args,
-                required = [(iter, Some(cloned_iter_type.clone()))]
-            );
-            // Return self:
-            Ok(iter.clone())
-        }
-    };
-
-    extend_class!(context, iter_type, {
-        "__iter__" => context.new_rustfunc(iter_func)
-    });
-}
-
-// Sequence iterator:
-fn iter_new(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(vm, args, required = [(iter_target, None)]);
+type PyIteratorValueRef = PyRef<PyIteratorValue>;
 
-    get_iter(vm, iter_target)
-}
-
-fn iter_next(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(vm, args, required = [(iter, Some(vm.ctx.iter_type()))]);
-
-    if let Some(PyIteratorValue {
-        ref position,
-        iterated_obj: ref iterated_obj_ref,
-    }) = iter.payload()
-    {
+impl PyIteratorValueRef {
+    fn next(self, vm: &VirtualMachine) -> PyResult {
+        let position = &self.position;
+        let iterated_obj_ref = &self.iterated_obj;
         if let Some(range) = iterated_obj_ref.payload::<PyRange>() {
             if let Some(int) = range.get(position.get()) {
                 position.set(position.get() + 1);
@@ -141,8 +110,10 @@ fn iter_next(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
                 Err(new_stop_iteration(vm))
             }
         }
-    } else {
-        panic!("NOT IMPL");
+    }
+
+    fn iter(self, _vm: &VirtualMachine) -> Self {
+        self
     }
 }
 
@@ -155,10 +126,9 @@ pub fn init(context: &PyContext) {
                     supply its own iterator, or be a sequence.\n\
                     In the second form, the callable is called until it returns the sentinel.";
 
-    iter_type_init(context, iter_type);
     extend_class!(context, iter_type, {
-        "__new__" => context.new_rustfunc(iter_new),
-        "__next__" => context.new_rustfunc(iter_next),
-        "__doc__" => context.new_str(iter_doc.to_string())
+        "__next__" => context.new_rustfunc(PyIteratorValueRef::next),
+        "__iter__" => context.new_rustfunc(PyIteratorValueRef::iter),
+        "__doc__" => context.new_str(iter_doc.to_string()),
     });
 }

vm/src/obj/objmap.rs

@@ -46,6 +46,10 @@ impl PyMapRef {
         // the mapper itself can raise StopIteration which does stop the map iteration
         vm.invoke(self.mapper.clone(), next_objs)
     }
+
+    fn iter(self, _vm: &VirtualMachine) -> Self {
+        self
+    }
 }
 
 pub fn init(context: &PyContext) {
@@ -55,10 +59,10 @@ pub fn init(context: &PyContext) {
                    Make an iterator that computes the function using arguments from\n\
                    each of the iterables.  Stops when the shortest iterable is exhausted.";
 
-    objiter::iter_type_init(context, map_type);
     extend_class!(context, map_type, {
         "__new__" => context.new_rustfunc(map_new),
         "__next__" => context.new_rustfunc(PyMapRef::next),
+        "__iter__" => context.new_rustfunc(PyMapRef::iter),
         "__doc__" => context.new_str(map_doc.to_string())
     });
 }

vm/src/obj/objzip.rs

@@ -40,13 +40,17 @@ impl PyZipRef {
             Ok(vm.ctx.new_tuple(next_objs))
         }
     }
+
+    fn iter(self, _vm: &VirtualMachine) -> Self {
+        self
+    }
 }
 
 pub fn init(context: &PyContext) {
     let zip_type = &context.zip_type;
-    objiter::iter_type_init(context, zip_type);
     extend_class!(context, zip_type, {
         "__new__" => context.new_rustfunc(zip_new),
-        "__next__" => context.new_rustfunc(PyZipRef::next)
+        "__next__" => context.new_rustfunc(PyZipRef::next),
+        "__iter__" => context.new_rustfunc(PyZipRef::iter),
     });
 }