Skip to content

Commit d5e2b3a

Browse files
adrian17adrian17
committed
Convert iterators next() to new args style
Except objiter, as it'll need to be reworked anyway
1 parent 149aefe commit d5e2b3a

File tree

4 files changed

+29
-56
lines changed

4 files changed

+29
-56
lines changed

‎vm/src/obj/objenumerate.rs

Lines changed: 7 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@ use std::ops::AddAssign;
44
use num_bigint::BigInt;
55
use num_traits::Zero;
66

7-
use crate::function::{OptionalArg, PyFuncArgs};
8-
use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol};
7+
use crate::function::OptionalArg;
8+
use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue};
99
use crate::vm::VirtualMachine;
1010

1111
use super::objint::PyIntRef;
@@ -44,18 +44,10 @@ fn enumerate_new(
4444
.into_ref_with_type(vm, cls)
4545
}
4646

47-
fn enumerate_next(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
48-
arg_check!(
49-
vm,
50-
args,
51-
required = [(enumerate, Some(vm.ctx.enumerate_type()))]
52-
);
53-
54-
if let Some(PyEnumerate {
55-
ref counter,
56-
ref iterator,
57-
}) = enumerate.payload()
58-
{
47+
impl PyEnumerateRef {
48+
fn next(self, vm: &VirtualMachine) -> PyResult {
49+
let iterator = &self.iterator;
50+
let counter = &self.counter;
5951
let next_obj = objiter::call_next(vm, iterator)?;
6052
let result = vm
6153
.ctx
@@ -64,8 +56,6 @@ fn enumerate_next(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
6456
AddAssign::add_assign(&mut counter.borrow_mut() as &mut BigInt, 1);
6557

6658
Ok(result)
67-
} else {
68-
panic!("enumerate doesn't have correct payload");
6959
}
7060
}
7161

@@ -74,6 +64,6 @@ pub fn init(context: &PyContext) {
7464
objiter::iter_type_init(context, enumerate_type);
7565
extend_class!(context, enumerate_type, {
7666
"__new__" => context.new_rustfunc(enumerate_new),
77-
"__next__" => context.new_rustfunc(enumerate_next)
67+
"__next__" => context.new_rustfunc(PyEnumerateRef::next)
7868
});
7969
}

‎vm/src/obj/objfilter.rs

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
1-
use crate::function::PyFuncArgs;
2-
use crate::pyobject::{IdProtocol, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol};
1+
use crate::pyobject::{IdProtocol, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
32
use crate::vm::VirtualMachine; // Required for arg_check! to use isinstance
43

54
use super::objbool;
@@ -35,14 +34,10 @@ fn filter_new(
3534
.into_ref_with_type(vm, cls)
3635
}
3736

38-
fn filter_next(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
39-
arg_check!(vm, args, required = [(filter, Some(vm.ctx.filter_type()))]);
40-
41-
if let Some(PyFilter {
42-
ref predicate,
43-
ref iterator,
44-
}) = filter.payload()
45-
{
37+
impl PyFilterRef {
38+
fn next(self, vm: &VirtualMachine) -> PyResult {
39+
let predicate = &self.predicate;
40+
let iterator = &self.iterator;
4641
loop {
4742
let next_obj = objiter::call_next(vm, iterator)?;
4843
let predicate_value = if predicate.is(&vm.get_none()) {
@@ -56,8 +51,6 @@ fn filter_next(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
5651
return Ok(next_obj);
5752
}
5853
}
59-
} else {
60-
panic!("filter doesn't have correct payload");
6154
}
6255
}
6356

@@ -74,6 +67,6 @@ pub fn init(context: &PyContext) {
7467
extend_class!(context, filter_type, {
7568
"__new__" => context.new_rustfunc(filter_new),
7669
"__doc__" => context.new_str(filter_doc.to_string()),
77-
"__next__" => context.new_rustfunc(filter_next)
70+
"__next__" => context.new_rustfunc(PyFilterRef::next)
7871
});
7972
}

‎vm/src/obj/objmap.rs

Lines changed: 8 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
use crate::function::{Args, PyFuncArgs};
2-
use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol};
1+
use crate::function::Args;
2+
use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue};
33
use crate::vm::VirtualMachine;
44

55
use super::objiter;
@@ -35,23 +35,16 @@ fn map_new(
3535
.into_ref_with_type(vm, cls.clone())
3636
}
3737

38-
fn map_next(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
39-
arg_check!(vm, args, required = [(map, Some(vm.ctx.map_type()))]);
40-
41-
if let Some(PyMap {
42-
ref mapper,
43-
ref iterators,
44-
}) = map.payload()
45-
{
46-
let next_objs = iterators
38+
impl PyMapRef {
39+
fn next(self, vm: &VirtualMachine) -> PyResult {
40+
let next_objs = self
41+
.iterators
4742
.iter()
4843
.map(|iterator| objiter::call_next(vm, iterator))
4944
.collect::<Result<Vec<_>, _>>()?;
5045

5146
// the mapper itself can raise StopIteration which does stop the map iteration
52-
vm.invoke(mapper.clone(), next_objs)
53-
} else {
54-
panic!("map doesn't have correct payload");
47+
vm.invoke(self.mapper.clone(), next_objs)
5548
}
5649
}
5750

@@ -65,7 +58,7 @@ pub fn init(context: &PyContext) {
6558
objiter::iter_type_init(context, map_type);
6659
extend_class!(context, map_type, {
6760
"__new__" => context.new_rustfunc(map_new),
68-
"__next__" => context.new_rustfunc(map_next),
61+
"__next__" => context.new_rustfunc(PyMapRef::next),
6962
"__doc__" => context.new_str(map_doc.to_string())
7063
});
7164
}

‎vm/src/obj/objzip.rs

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
use crate::function::{Args, PyFuncArgs};
2-
use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol};
1+
use crate::function::Args;
2+
use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue};
33
use crate::vm::VirtualMachine;
44

55
use super::objiter;
@@ -26,22 +26,19 @@ fn zip_new(cls: PyClassRef, iterables: Args, vm: &VirtualMachine) -> PyResult<Py
2626
PyZip { iterators }.into_ref_with_type(vm, cls)
2727
}
2828

29-
fn zip_next(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
30-
arg_check!(vm, args, required = [(zip, Some(vm.ctx.zip_type()))]);
31-
32-
if let Some(PyZip { ref iterators }) = zip.payload() {
33-
if iterators.is_empty() {
29+
impl PyZipRef {
30+
fn next(self, vm: &VirtualMachine) -> PyResult {
31+
if self.iterators.is_empty() {
3432
Err(objiter::new_stop_iteration(vm))
3533
} else {
36-
let next_objs = iterators
34+
let next_objs = self
35+
.iterators
3736
.iter()
3837
.map(|iterator| objiter::call_next(vm, iterator))
3938
.collect::<Result<Vec<_>, _>>()?;
4039

4140
Ok(vm.ctx.new_tuple(next_objs))
4241
}
43-
} else {
44-
panic!("zip doesn't have correct payload");
4542
}
4643
}
4744

@@ -50,6 +47,6 @@ pub fn init(context: &PyContext) {
5047
objiter::iter_type_init(context, zip_type);
5148
extend_class!(context, zip_type, {
5249
"__new__" => context.new_rustfunc(zip_new),
53-
"__next__" => context.new_rustfunc(zip_next)
50+
"__next__" => context.new_rustfunc(PyZipRef::next)
5451
});
5552
}

0 commit comments

Comments
 (0)