Skip to content

Commit 09b99fd

Browse files
youknowoneyouknowone
committed
PyIter protocol
1 parent db4e42f commit 09b99fd

18 files changed

+270
-242
lines changed

‎vm/src/builtins/dict.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -87,13 +87,13 @@ impl PyDict {
8787
return Err(vm.new_runtime_error("dict mutated during update".to_owned()));
8888
}
8989
} else if let Some(keys) = vm.get_method(dict_obj.clone(), "keys") {
90-
let keys = iterator::get_iter(vm, vm.invoke(&keys?, ())?)?;
90+
let keys = vm.invoke(&keys?, ())?.get_iter(vm)?;
9191
while let Some(key) = iterator::get_next_object(vm, &keys)? {
9292
let val = dict_obj.get_item(key.clone(), vm)?;
9393
dict.insert(vm, key, val)?;
9494
}
9595
} else {
96-
let iter = iterator::get_iter(vm, dict_obj)?;
96+
let iter = dict_obj.get_iter(vm)?;
9797
loop {
9898
fn err(vm: &VirtualMachine) -> PyBaseExceptionRef {
9999
vm.new_value_error("Iterator must have exactly two elements".to_owned())
@@ -102,7 +102,7 @@ impl PyDict {
102102
Some(obj) => obj,
103103
None => break,
104104
};
105-
let elem_iter = iterator::get_iter(vm, element)?;
105+
let elem_iter = element.get_iter(vm)?;
106106
let key = iterator::get_next_object(vm, &elem_iter)?.ok_or_else(|| err(vm))?;
107107
let value =
108108
iterator::get_next_object(vm, &elem_iter)?.ok_or_else(|| err(vm))?;

‎vm/src/builtins/enumerate.rs

+11-9
Original file line numberDiff line numberDiff line change
@@ -11,16 +11,17 @@ use super::iter::{
1111
};
1212
use super::pytype::PyTypeRef;
1313
use crate::function::OptionalArg;
14+
use crate::protocol::PyIter;
1415
use crate::slots::{IteratorIterable, SlotConstructor, SlotIterator};
1516
use crate::vm::VirtualMachine;
16-
use crate::{iterator, ItemProtocol, TypeProtocol};
1717
use crate::{IntoPyObject, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
18+
use crate::{ItemProtocol, TypeProtocol};
1819

1920
#[pyclass(module = false, name = "enumerate")]
2021
#[derive(Debug)]
2122
pub struct PyEnumerate {
2223
counter: PyRwLock<BigInt>,
23-
iterator: PyObjectRef,
24+
iterator: PyIter,
2425
}
2526

2627
impl PyValue for PyEnumerate {
@@ -32,19 +33,20 @@ impl PyValue for PyEnumerate {
3233
#[derive(FromArgs)]
3334
pub struct EnumerateArgs {
3435
#[pyarg(any)]
35-
iterable: PyObjectRef,
36+
iterator: PyIter,
3637
#[pyarg(any, optional)]
3738
start: OptionalArg<PyIntRef>,
3839
}
3940

4041
impl SlotConstructor for PyEnumerate {
4142
type Args = EnumerateArgs;
4243

43-
fn py_new(cls: PyTypeRef, args: Self::Args, vm: &VirtualMachine) -> PyResult {
44-
let counter = args
45-
.start
46-
.map_or_else(BigInt::zero, |start| start.as_bigint().clone());
47-
let iterator = iterator::get_iter(vm, args.iterable)?;
44+
fn py_new(
45+
cls: PyTypeRef,
46+
Self::Args { iterator, start }: Self::Args,
47+
vm: &VirtualMachine,
48+
) -> PyResult {
49+
let counter = start.map_or_else(BigInt::zero, |start| start.as_bigint().clone());
4850
PyEnumerate {
4951
counter: PyRwLock::new(counter),
5052
iterator,
@@ -59,7 +61,7 @@ impl PyEnumerate {}
5961
impl IteratorIterable for PyEnumerate {}
6062
impl SlotIterator for PyEnumerate {
6163
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
62-
let next_obj = iterator::call_next(vm, &zelf.iterator)?;
64+
let next_obj = zelf.iterator.next(vm)?;
6365
let mut counter = zelf.counter.write();
6466
let position = counter.clone();
6567
*counter += 1;

‎vm/src/builtins/filter.rs

+5-20
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
use super::pytype::PyTypeRef;
2-
use crate::iterator;
2+
use crate::protocol::PyIter;
33
use crate::slots::{IteratorIterable, SlotConstructor, SlotIterator};
44
use crate::vm::VirtualMachine;
55
use crate::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
@@ -12,7 +12,7 @@ use crate::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
1212
#[derive(Debug)]
1313
pub struct PyFilter {
1414
predicate: PyObjectRef,
15-
iterator: PyObjectRef,
15+
iterator: PyIter,
1616
}
1717

1818
impl PyValue for PyFilter {
@@ -21,24 +21,10 @@ impl PyValue for PyFilter {
2121
}
2222
}
2323

24-
#[derive(FromArgs)]
25-
pub struct FilterArgs {
26-
#[pyarg(positional)]
27-
function: PyObjectRef,
28-
#[pyarg(positional)]
29-
iterable: PyObjectRef,
30-
}
31-
3224
impl SlotConstructor for PyFilter {
33-
type Args = FilterArgs;
34-
35-
fn py_new(
36-
cls: PyTypeRef,
37-
Self::Args { function, iterable }: Self::Args,
38-
vm: &VirtualMachine,
39-
) -> PyResult {
40-
let iterator = iterator::get_iter(vm, iterable)?;
25+
type Args = (PyObjectRef, PyIter);
4126

27+
fn py_new(cls: PyTypeRef, (function, iterator): Self::Args, vm: &VirtualMachine) -> PyResult {
4228
Self {
4329
predicate: function,
4430
iterator,
@@ -54,9 +40,8 @@ impl IteratorIterable for PyFilter {}
5440
impl SlotIterator for PyFilter {
5541
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
5642
let predicate = &zelf.predicate;
57-
let iterator = &zelf.iterator;
5843
loop {
59-
let next_obj = iterator::call_next(vm, iterator)?;
44+
let next_obj = zelf.iterator.next(vm)?;
6045
let predicate_value = if vm.is_none(predicate) {
6146
next_obj.clone()
6247
} else {

‎vm/src/builtins/make_module.rs

+7-6
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ mod decl {
2828
use crate::function::{
2929
ArgCallable, ArgIterable, Args, FuncArgs, KwArgs, OptionalArg, OptionalOption,
3030
};
31-
use crate::iterator;
31+
use crate::protocol::PyIter;
3232
use crate::readline::{Readline, ReadlineResult};
3333
use crate::scope::Scope;
3434
use crate::slots::PyComparisonOp;
@@ -421,14 +421,15 @@ mod decl {
421421
iter_target: PyObjectRef,
422422
sentinel: OptionalArg<PyObjectRef>,
423423
vm: &VirtualMachine,
424-
) -> PyResult {
424+
) -> PyResult<PyIter> {
425425
if let OptionalArg::Present(sentinel) = sentinel {
426426
let callable = ArgCallable::try_from_object(vm, iter_target)?;
427-
Ok(PyCallableIterator::new(callable, sentinel)
427+
let iterator = PyCallableIterator::new(callable, sentinel)
428428
.into_ref(vm)
429-
.into_object())
429+
.into_object();
430+
Ok(PyIter::new(iterator))
430431
} else {
431-
iterator::get_iter(vm, iter_target)
432+
iter_target.get_iter(vm)
432433
}
433434
}
434435

@@ -519,7 +520,7 @@ mod decl {
519520
default_value: OptionalArg<PyObjectRef>,
520521
vm: &VirtualMachine,
521522
) -> PyResult {
522-
iterator::call_next(vm, &iterator).or_else(|err| {
523+
PyIter::new(iterator).next(vm).or_else(|err| {
523524
if err.isinstance(&vm.ctx.exceptions.stop_iteration) {
524525
default_value.ok_or(err)
525526
} else {

‎vm/src/builtins/map.rs

+8-14
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use super::pytype::PyTypeRef;
22
use crate::function::Args;
33
use crate::iterator;
4+
use crate::protocol::PyIter;
45
use crate::slots::{IteratorIterable, SlotConstructor, SlotIterator};
56
use crate::vm::VirtualMachine;
67
use crate::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
@@ -13,7 +14,7 @@ use crate::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
1314
#[derive(Debug)]
1415
pub struct PyMap {
1516
mapper: PyObjectRef,
16-
iterators: Vec<PyObjectRef>,
17+
iterators: Vec<PyIter>,
1718
}
1819

1920
impl PyValue for PyMap {
@@ -23,18 +24,11 @@ impl PyValue for PyMap {
2324
}
2425

2526
impl SlotConstructor for PyMap {
26-
type Args = (PyObjectRef, Args<PyObjectRef>);
27+
type Args = (PyObjectRef, Args<PyIter>);
2728

28-
fn py_new(cls: PyTypeRef, (function, iterables): Self::Args, vm: &VirtualMachine) -> PyResult {
29-
let iterators = iterables
30-
.into_iter()
31-
.map(|iterable| iterator::get_iter(vm, iterable))
32-
.collect::<Result<Vec<_>, _>>()?;
33-
PyMap {
34-
mapper: function,
35-
iterators,
36-
}
37-
.into_pyresult_with_type(vm, cls)
29+
fn py_new(cls: PyTypeRef, (mapper, iterators): Self::Args, vm: &VirtualMachine) -> PyResult {
30+
let iterators = iterators.into_vec();
31+
PyMap { mapper, iterators }.into_pyresult_with_type(vm, cls)
3832
}
3933
}
4034

@@ -43,7 +37,7 @@ impl PyMap {
4337
#[pymethod(magic)]
4438
fn length_hint(&self, vm: &VirtualMachine) -> PyResult<usize> {
4539
self.iterators.iter().try_fold(0, |prev, cur| {
46-
let cur = iterator::length_hint(vm, cur.clone())?.unwrap_or(0);
40+
let cur = iterator::length_hint(vm, cur.as_object().clone())?.unwrap_or(0);
4741
let max = std::cmp::max(prev, cur);
4842
Ok(max)
4943
})
@@ -56,7 +50,7 @@ impl SlotIterator for PyMap {
5650
let next_objs = zelf
5751
.iterators
5852
.iter()
59-
.map(|iterator| iterator::call_next(vm, iterator))
53+
.map(|iterator| iterator.next(vm))
6054
.collect::<Result<Vec<_>, _>>()?;
6155

6256
// the mapper itself can raise StopIteration which does stop the map iteration

‎vm/src/builtins/mappingproxy.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ use super::dict::PyDict;
22
use super::pystr::PyStrRef;
33
use super::pytype::PyTypeRef;
44
use crate::function::OptionalArg;
5-
use crate::iterator;
65
use crate::slots::{Iterable, SlotConstructor};
76
use crate::vm::VirtualMachine;
87
use crate::{
@@ -140,7 +139,8 @@ impl Iterable for PyMappingProxy {
140139
PyDict::from_attributes(c.attributes.read().clone(), vm)?.into_pyobject(vm)
141140
}
142141
};
143-
iterator::get_iter(vm, obj)
142+
let iter = obj.get_iter(vm)?;
143+
Ok(iter.into_object())
144144
}
145145
}
146146

‎vm/src/builtins/range.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ fn iter_search(
3535
vm: &VirtualMachine,
3636
) -> PyResult<usize> {
3737
let mut count = 0;
38-
let iter = iterator::get_iter(vm, obj)?;
38+
let iter = obj.get_iter(vm)?;
3939
while let Some(element) = iterator::get_next_object(vm, &iter)? {
4040
if vm.bool_eq(&item, &element)? {
4141
match flag {

‎vm/src/builtins/zip.rs

+7-10
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,14 @@
11
use super::pytype::PyTypeRef;
22
use crate::function::Args;
3-
use crate::iterator;
3+
use crate::protocol::PyIter;
44
use crate::slots::{IteratorIterable, SlotConstructor, SlotIterator};
55
use crate::vm::VirtualMachine;
6-
use crate::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
6+
use crate::{PyClassImpl, PyContext, PyRef, PyResult, PyValue};
77

88
#[pyclass(module = false, name = "zip")]
99
#[derive(Debug)]
1010
pub struct PyZip {
11-
iterators: Vec<PyObjectRef>,
11+
iterators: Vec<PyIter>,
1212
}
1313

1414
impl PyValue for PyZip {
@@ -18,13 +18,10 @@ impl PyValue for PyZip {
1818
}
1919

2020
impl SlotConstructor for PyZip {
21-
type Args = Args;
21+
type Args = Args<PyIter>;
2222

23-
fn py_new(cls: PyTypeRef, iterables: Self::Args, vm: &VirtualMachine) -> PyResult {
24-
let iterators = iterables
25-
.into_iter()
26-
.map(|iterable| iterator::get_iter(vm, iterable))
27-
.collect::<Result<Vec<_>, _>>()?;
23+
fn py_new(cls: PyTypeRef, iterators: Self::Args, vm: &VirtualMachine) -> PyResult {
24+
let iterators = iterators.into_vec();
2825
PyZip { iterators }.into_pyresult_with_type(vm, cls)
2926
}
3027
}
@@ -41,7 +38,7 @@ impl SlotIterator for PyZip {
4138
let next_objs = zelf
4239
.iterators
4340
.iter()
44-
.map(|iterator| iterator::call_next(vm, iterator))
41+
.map(|iterator| iterator.next(vm))
4542
.collect::<Result<Vec<_>, _>>()?;
4643

4744
Ok(vm.ctx.new_tuple(next_objs))

‎vm/src/frame.rs

+5-4
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ use crate::coroutine::Coro;
2525
use crate::exceptions::{self, ExceptionCtor, PyBaseExceptionRef};
2626
use crate::function::FuncArgs;
2727
use crate::iterator;
28+
use crate::protocol::PyIter;
2829
use crate::scope::Scope;
2930
use crate::slots::PyComparisonOp;
3031
use crate::vm::VirtualMachine;
@@ -862,8 +863,8 @@ impl ExecutingFrame<'_> {
862863
}
863864
bytecode::Instruction::GetIter => {
864865
let iterated_obj = self.pop_value();
865-
let iter_obj = iterator::get_iter(vm, iterated_obj)?;
866-
self.push_value(iter_obj);
866+
let iter_obj = iterated_obj.get_iter(vm)?;
867+
self.push_value(iter_obj.into_object());
867868
Ok(None)
868869
}
869870
bytecode::Instruction::GetAwaitable => {
@@ -1447,7 +1448,7 @@ impl ExecutingFrame<'_> {
14471448
fn _send(&self, coro: &PyObjectRef, val: PyObjectRef, vm: &VirtualMachine) -> PyResult {
14481449
match self.builtin_coro(coro) {
14491450
Some(coro) => coro.send(val, vm),
1450-
None if vm.is_none(&val) => iterator::call_next(vm, coro),
1451+
None if vm.is_none(&val) => PyIter::new(coro).next(vm),
14511452
None => {
14521453
let meth = vm.get_attribute(coro.clone(), "send")?;
14531454
vm.invoke(&meth, (val,))
@@ -1517,7 +1518,7 @@ impl ExecutingFrame<'_> {
15171518

15181519
/// The top of stack contains the iterator, lets push it forward
15191520
fn execute_for_iter(&mut self, vm: &VirtualMachine, target: bytecode::Label) -> FrameResult {
1520-
let top_of_stack = self.last_value();
1521+
let top_of_stack = PyIter::new(self.last_value());
15211522
let next_obj = iterator::get_next_object(vm, &top_of_stack);
15221523

15231524
// Check the next object:

‎vm/src/function/argument.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use super::IntoFuncArgs;
2-
use crate::builtins::iter::PySequenceIterator;
32
use crate::{
4-
iterator, PyObjectRef, PyResult, PyValue, TryFromObject, TypeProtocol, VirtualMachine,
3+
builtins::iter::PySequenceIterator, iterator, protocol::PyIter, PyObjectRef, PyResult, PyValue,
4+
TryFromObject, TypeProtocol, VirtualMachine,
55
};
66
use std::marker::PhantomData;
77

@@ -102,7 +102,7 @@ where
102102
type Item = PyResult<T>;
103103

104104
fn next(&mut self) -> Option<Self::Item> {
105-
iterator::get_next_object(self.vm, &self.obj)
105+
iterator::get_next_object(self.vm, &PyIter::new(&self.obj))
106106
.transpose()
107107
.map(|x| x.and_then(|obj| T::try_from_object(self.vm, obj)))
108108
}

0 commit comments

Comments
 (0)