RustPython · coolreader18 · Mar 28, 2019 · Mar 28, 2019 · Mar 28, 2019
diff --git a/vm/src/obj/objenumerate.rs b/vm/src/obj/objenumerate.rs
@@ -4,8 +4,8 @@ use std::ops::AddAssign;
 use num_bigint::BigInt;
 use num_traits::Zero;
 
-use crate::function::{OptionalArg, PyFuncArgs};
-use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol};
+use crate::function::OptionalArg;
+use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue};
 use crate::vm::VirtualMachine;
 
 use super::objint::PyIntRef;
@@ -44,18 +44,10 @@ fn enumerate_new(
     .into_ref_with_type(vm, cls)
 }
 
-fn enumerate_next(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(
-        vm,
-        args,
-        required = [(enumerate, Some(vm.ctx.enumerate_type()))]
-    );
-
-    if let Some(PyEnumerate {
-        ref counter,
-        ref iterator,
-    }) = enumerate.payload()
-    {
+impl PyEnumerateRef {
+    fn next(self, vm: &VirtualMachine) -> PyResult {
+        let iterator = &self.iterator;
+        let counter = &self.counter;
         let next_obj = objiter::call_next(vm, iterator)?;
         let result = vm
             .ctx
@@ -64,8 +56,6 @@ fn enumerate_next(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
         AddAssign::add_assign(&mut counter.borrow_mut() as &mut BigInt, 1);
 
         Ok(result)
-    } else {
-        panic!("enumerate doesn't have correct payload");
     }
 }
 
@@ -74,6 +64,6 @@ pub fn init(context: &PyContext) {
     objiter::iter_type_init(context, enumerate_type);
     extend_class!(context, enumerate_type, {
         "__new__" => context.new_rustfunc(enumerate_new),
-        "__next__" => context.new_rustfunc(enumerate_next)
+        "__next__" => context.new_rustfunc(PyEnumerateRef::next)
     });
 }
diff --git a/vm/src/obj/objfilter.rs b/vm/src/obj/objfilter.rs
@@ -1,5 +1,4 @@
-use crate::function::PyFuncArgs;
-use crate::pyobject::{IdProtocol, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol};
+use crate::pyobject::{IdProtocol, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
 use crate::vm::VirtualMachine; // Required for arg_check! to use isinstance
 
 use super::objbool;
@@ -35,14 +34,10 @@ fn filter_new(
     .into_ref_with_type(vm, cls)
 }
 
-fn filter_next(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(vm, args, required = [(filter, Some(vm.ctx.filter_type()))]);
-
-    if let Some(PyFilter {
-        ref predicate,
-        ref iterator,
-    }) = filter.payload()
-    {
+impl PyFilterRef {
+    fn next(self, vm: &VirtualMachine) -> PyResult {
+        let predicate = &self.predicate;
+        let iterator = &self.iterator;
         loop {
             let next_obj = objiter::call_next(vm, iterator)?;
             let predicate_value = if predicate.is(&vm.get_none()) {
@@ -56,8 +51,6 @@ fn filter_next(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
                 return Ok(next_obj);
             }
         }
-    } else {
-        panic!("filter doesn't have correct payload");
     }
 }
 
@@ -74,6 +67,6 @@ 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(filter_next)
+        "__next__" => context.new_rustfunc(PyFilterRef::next)
     });
 }
diff --git a/vm/src/obj/objmap.rs b/vm/src/obj/objmap.rs
@@ -1,5 +1,5 @@
-use crate::function::{Args, PyFuncArgs};
-use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol};
+use crate::function::Args;
+use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue};
 use crate::vm::VirtualMachine;
 
 use super::objiter;
@@ -35,23 +35,16 @@ fn map_new(
     .into_ref_with_type(vm, cls.clone())
 }
 
-fn map_next(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(vm, args, required = [(map, Some(vm.ctx.map_type()))]);
-
-    if let Some(PyMap {
-        ref mapper,
-        ref iterators,
-    }) = map.payload()
-    {
-        let next_objs = iterators
+impl PyMapRef {
+    fn next(self, vm: &VirtualMachine) -> PyResult {
+        let next_objs = self
+            .iterators
             .iter()
             .map(|iterator| objiter::call_next(vm, iterator))
             .collect::<Result<Vec<_>, _>>()?;
 
         // the mapper itself can raise StopIteration which does stop the map iteration
-        vm.invoke(mapper.clone(), next_objs)
-    } else {
-        panic!("map doesn't have correct payload");
+        vm.invoke(self.mapper.clone(), next_objs)
     }
 }
 
@@ -65,7 +58,7 @@ pub fn init(context: &PyContext) {
     objiter::iter_type_init(context, map_type);
     extend_class!(context, map_type, {
         "__new__" => context.new_rustfunc(map_new),
-        "__next__" => context.new_rustfunc(map_next),
+        "__next__" => context.new_rustfunc(PyMapRef::next),
         "__doc__" => context.new_str(map_doc.to_string())
     });
 }
diff --git a/vm/src/obj/objslice.rs b/vm/src/obj/objslice.rs
@@ -66,38 +66,25 @@ fn slice_new(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     .map(|x| x.into_object())
 }
 
-fn get_property_value(vm: &VirtualMachine, value: &Option<BigInt>) -> PyResult {
+fn get_property_value(vm: &VirtualMachine, value: &Option<BigInt>) -> PyObjectRef {
     if let Some(value) = value {
-        Ok(vm.ctx.new_int(value.clone()))
+        vm.ctx.new_int(value.clone())
     } else {
-        Ok(vm.get_none())
+        vm.get_none()
     }
 }
 
-fn slice_start(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(vm, args, required = [(slice, Some(vm.ctx.slice_type()))]);
-    if let Some(PySlice { start, .. }) = &slice.payload() {
-        get_property_value(vm, start)
-    } else {
-        panic!("Slice has incorrect payload.");
+impl PySliceRef {
+    fn start(self, vm: &VirtualMachine) -> PyObjectRef {
+        get_property_value(vm, &self.start)
     }
-}
 
-fn slice_stop(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(vm, args, required = [(slice, Some(vm.ctx.slice_type()))]);
-    if let Some(PySlice { stop, .. }) = &slice.payload() {
-        get_property_value(vm, stop)
-    } else {
-        panic!("Slice has incorrect payload.");
+    fn stop(self, vm: &VirtualMachine) -> PyObjectRef {
+        get_property_value(vm, &self.stop)
     }
-}
 
-fn slice_step(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(vm, args, required = [(slice, Some(vm.ctx.slice_type()))]);
-    if let Some(PySlice { step, .. }) = &slice.payload() {
-        get_property_value(vm, step)
-    } else {
-        panic!("Slice has incorrect payload.");
+    fn step(self, vm: &VirtualMachine) -> PyObjectRef {
+        get_property_value(vm, &self.step)
     }
 }
 
@@ -106,8 +93,8 @@ pub fn init(context: &PyContext) {
 
     extend_class!(context, slice_type, {
         "__new__" => context.new_rustfunc(slice_new),
-        "start" => context.new_property(slice_start),
-        "stop" => context.new_property(slice_stop),
-        "step" => context.new_property(slice_step)
+        "start" => context.new_property(PySliceRef::start),
+        "stop" => context.new_property(PySliceRef::stop),
+        "step" => context.new_property(PySliceRef::step)
     });
 }
diff --git a/vm/src/obj/objzip.rs b/vm/src/obj/objzip.rs
@@ -1,5 +1,5 @@
-use crate::function::{Args, PyFuncArgs};
-use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol};
+use crate::function::Args;
+use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue};
 use crate::vm::VirtualMachine;
 
 use super::objiter;
@@ -26,22 +26,19 @@ fn zip_new(cls: PyClassRef, iterables: Args, vm: &VirtualMachine) -> PyResult<Py
     PyZip { iterators }.into_ref_with_type(vm, cls)
 }
 
-fn zip_next(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(vm, args, required = [(zip, Some(vm.ctx.zip_type()))]);
-
-    if let Some(PyZip { ref iterators }) = zip.payload() {
-        if iterators.is_empty() {
+impl PyZipRef {
+    fn next(self, vm: &VirtualMachine) -> PyResult {
+        if self.iterators.is_empty() {
             Err(objiter::new_stop_iteration(vm))
         } else {
-            let next_objs = iterators
+            let next_objs = self
+                .iterators
                 .iter()
                 .map(|iterator| objiter::call_next(vm, iterator))
                 .collect::<Result<Vec<_>, _>>()?;
 
             Ok(vm.ctx.new_tuple(next_objs))
         }
-    } else {
-        panic!("zip doesn't have correct payload");
     }
 }
 
@@ -50,6 +47,6 @@ pub fn init(context: &PyContext) {
     objiter::iter_type_init(context, zip_type);
     extend_class!(context, zip_type, {
         "__new__" => context.new_rustfunc(zip_new),
-        "__next__" => context.new_rustfunc(zip_next)
+        "__next__" => context.new_rustfunc(PyZipRef::next)
     });
 }