diff --git a/vm/src/obj/objenumerate.rs b/vm/src/obj/objenumerate.rs index 68d1e49f6b..85781bb7a0 100644 --- 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 index e547d2c43e..07c512f6dd 100644 --- 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 index 49cc539e58..fbf0fa035a 100644 --- 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::, _>>()?; // 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 index 3d955b3f2c..c82a3ca6ef 100644 --- 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) -> PyResult { +fn get_property_value(vm: &VirtualMachine, value: &Option) -> 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 index e628cb218a..a12152eeec 100644 --- 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 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::, _>>()?; 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) }); }