Skip to content

Commit ed9b172

Browse files
youknowoneyouknowone
committed
PyIter protocol
1 parent 8a46e63 commit ed9b172

File tree

18 files changed

+275
-245
lines changed

18 files changed

+275
-245
lines changed

‎vm/src/builtins/dict.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -85,13 +85,13 @@ impl PyDict {
8585
return Err(vm.new_runtime_error("dict mutated during update".to_owned()));
8686
}
8787
} else if let Some(keys) = vm.get_method(dict_obj.clone(), "keys") {
88-
let keys = iterator::get_iter(vm, vm.invoke(&keys?, ())?)?;
88+
let keys = vm.invoke(&keys?, ())?.get_iter(vm)?;
8989
while let Some(key) = iterator::get_next_object(vm, &keys)? {
9090
let val = dict_obj.get_item(key.clone(), vm)?;
9191
dict.insert(vm, key, val)?;
9292
}
9393
} else {
94-
let iter = iterator::get_iter(vm, dict_obj)?;
94+
let iter = dict_obj.get_iter(vm)?;
9595
loop {
9696
fn err(vm: &VirtualMachine) -> PyBaseExceptionRef {
9797
vm.new_value_error("Iterator must have exactly two elements".to_owned())
@@ -100,7 +100,7 @@ impl PyDict {
100100
Some(obj) => obj,
101101
None => break,
102102
};
103-
let elem_iter = iterator::get_iter(vm, element)?;
103+
let elem_iter = element.get_iter(vm)?;
104104
let key = iterator::get_next_object(vm, &elem_iter)?.ok_or_else(|| err(vm))?;
105105
let value =
106106
iterator::get_next_object(vm, &elem_iter)?.ok_or_else(|| err(vm))?;

‎vm/src/builtins/enumerate.rs

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use super::{
66
use crate::common::lock::PyRwLock;
77
use crate::{
88
function::OptionalArg,
9-
iterator,
9+
protocol::PyIter,
1010
slots::{IteratorIterable, SlotConstructor, SlotIterator},
1111
IntoPyObject, ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue,
1212
TypeProtocol, VirtualMachine,
@@ -19,7 +19,7 @@ use num_traits::Zero;
1919
#[derive(Debug)]
2020
pub struct PyEnumerate {
2121
counter: PyRwLock<BigInt>,
22-
iterator: PyObjectRef,
22+
iterator: PyIter,
2323
}
2424

2525
impl PyValue for PyEnumerate {
@@ -30,19 +30,20 @@ impl PyValue for PyEnumerate {
3030

3131
#[derive(FromArgs)]
3232
pub struct EnumerateArgs {
33-
iterable: PyObjectRef,
33+
iterator: PyIter,
3434
#[pyarg(any, optional)]
3535
start: OptionalArg<PyIntRef>,
3636
}
3737

3838
impl SlotConstructor for PyEnumerate {
3939
type Args = EnumerateArgs;
4040

41-
fn py_new(cls: PyTypeRef, args: Self::Args, vm: &VirtualMachine) -> PyResult {
42-
let counter = args
43-
.start
44-
.map_or_else(BigInt::zero, |start| start.as_bigint().clone());
45-
let iterator = iterator::get_iter(vm, args.iterable)?;
41+
fn py_new(
42+
cls: PyTypeRef,
43+
Self::Args { iterator, start }: Self::Args,
44+
vm: &VirtualMachine,
45+
) -> PyResult {
46+
let counter = start.map_or_else(BigInt::zero, |start| start.as_bigint().clone());
4647
PyEnumerate {
4748
counter: PyRwLock::new(counter),
4849
iterator,
@@ -57,7 +58,7 @@ impl PyEnumerate {}
5758
impl IteratorIterable for PyEnumerate {}
5859
impl SlotIterator for PyEnumerate {
5960
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
60-
let next_obj = iterator::call_next(vm, &zelf.iterator)?;
61+
let next_obj = zelf.iterator.next(vm)?;
6162
let mut counter = zelf.counter.write();
6263
let position = counter.clone();
6364
*counter += 1;

‎vm/src/builtins/filter.rs

Lines changed: 5 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use super::PyTypeRef;
22
use crate::{
3-
iterator,
3+
protocol::PyIter,
44
slots::{IteratorIterable, SlotConstructor, SlotIterator},
55
PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, VirtualMachine,
66
};
@@ -13,7 +13,7 @@ use crate::{
1313
#[derive(Debug)]
1414
pub struct PyFilter {
1515
predicate: PyObjectRef,
16-
iterator: PyObjectRef,
16+
iterator: PyIter,
1717
}
1818

1919
impl PyValue for PyFilter {
@@ -22,24 +22,10 @@ impl PyValue for PyFilter {
2222
}
2323
}
2424

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

28+
fn py_new(cls: PyTypeRef, (function, iterator): Self::Args, vm: &VirtualMachine) -> PyResult {
4329
Self {
4430
predicate: function,
4531
iterator,
@@ -55,9 +41,8 @@ impl IteratorIterable for PyFilter {}
5541
impl SlotIterator for PyFilter {
5642
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
5743
let predicate = &zelf.predicate;
58-
let iterator = &zelf.iterator;
5944
loop {
60-
let next_obj = iterator::call_next(vm, iterator)?;
45+
let next_obj = zelf.iterator.next(vm)?;
6146
let predicate_value = if vm.is_none(predicate) {
6247
next_obj.clone()
6348
} else {

‎vm/src/builtins/make_module.rs

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,8 @@ mod decl {
2525
ArgBytesLike, ArgCallable, ArgIterable, FuncArgs, KwArgs, OptionalArg, OptionalOption,
2626
PosArgs,
2727
},
28-
iterator, py_io,
28+
protocol::PyIter,
29+
py_io,
2930
readline::{Readline, ReadlineResult},
3031
scope::Scope,
3132
slots::PyComparisonOp,
@@ -413,14 +414,15 @@ mod decl {
413414
iter_target: PyObjectRef,
414415
sentinel: OptionalArg<PyObjectRef>,
415416
vm: &VirtualMachine,
416-
) -> PyResult {
417+
) -> PyResult<PyIter> {
417418
if let OptionalArg::Present(sentinel) = sentinel {
418419
let callable = ArgCallable::try_from_object(vm, iter_target)?;
419-
Ok(PyCallableIterator::new(callable, sentinel)
420+
let iterator = PyCallableIterator::new(callable, sentinel)
420421
.into_ref(vm)
421-
.into_object())
422+
.into_object();
423+
Ok(PyIter::new(iterator))
422424
} else {
423-
iterator::get_iter(vm, iter_target)
425+
iter_target.get_iter(vm)
424426
}
425427
}
426428

@@ -511,7 +513,7 @@ mod decl {
511513
default_value: OptionalArg<PyObjectRef>,
512514
vm: &VirtualMachine,
513515
) -> PyResult {
514-
iterator::call_next(vm, &iterator).or_else(|err| {
516+
PyIter::new(iterator).next(vm).or_else(|err| {
515517
if err.isinstance(&vm.ctx.exceptions.stop_iteration) {
516518
default_value.ok_or(err)
517519
} else {

‎vm/src/builtins/map.rs

Lines changed: 8 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use super::PyTypeRef;
22
use crate::{
33
function::PosArgs,
44
iterator,
5+
protocol::PyIter,
56
slots::{IteratorIterable, SlotConstructor, SlotIterator},
67
PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, VirtualMachine,
78
};
@@ -14,7 +15,7 @@ use crate::{
1415
#[derive(Debug)]
1516
pub struct PyMap {
1617
mapper: PyObjectRef,
17-
iterators: Vec<PyObjectRef>,
18+
iterators: Vec<PyIter>,
1819
}
1920

2021
impl PyValue for PyMap {
@@ -24,18 +25,11 @@ impl PyValue for PyMap {
2425
}
2526

2627
impl SlotConstructor for PyMap {
27-
type Args = (PyObjectRef, PosArgs<PyObjectRef>);
28+
type Args = (PyObjectRef, PosArgs<PyIter>);
2829

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

@@ -44,7 +38,7 @@ impl PyMap {
4438
#[pymethod(magic)]
4539
fn length_hint(&self, vm: &VirtualMachine) -> PyResult<usize> {
4640
self.iterators.iter().try_fold(0, |prev, cur| {
47-
let cur = iterator::length_hint(vm, cur.clone())?.unwrap_or(0);
41+
let cur = iterator::length_hint(vm, cur.as_object().clone())?.unwrap_or(0);
4842
let max = std::cmp::max(prev, cur);
4943
Ok(max)
5044
})
@@ -57,7 +51,7 @@ impl SlotIterator for PyMap {
5751
let next_objs = zelf
5852
.iterators
5953
.iter()
60-
.map(|iterator| iterator::call_next(vm, iterator))
54+
.map(|iterator| iterator.next(vm))
6155
.collect::<Result<Vec<_>, _>>()?;
6256

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

‎vm/src/builtins/mappingproxy.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
use super::{PyDict, PyStrRef, PyTypeRef};
22
use crate::{
33
function::OptionalArg,
4-
iterator,
54
slots::{Iterable, SlotConstructor},
65
IntoPyObject, ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue,
76
TryFromObject, VirtualMachine,
@@ -137,7 +136,8 @@ impl Iterable for PyMappingProxy {
137136
PyDict::from_attributes(c.attributes.read().clone(), vm)?.into_pyobject(vm)
138137
}
139138
};
140-
iterator::get_iter(vm, obj)
139+
let iter = obj.get_iter(vm)?;
140+
Ok(iter.into_object())
141141
}
142142
}
143143

‎vm/src/builtins/range.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ fn iter_search(
3030
vm: &VirtualMachine,
3131
) -> PyResult<usize> {
3232
let mut count = 0;
33-
let iter = iterator::get_iter(vm, obj)?;
33+
let iter = obj.get_iter(vm)?;
3434
while let Some(element) = iterator::get_next_object(vm, &iter)? {
3535
if vm.bool_eq(&item, &element)? {
3636
match flag {

‎vm/src/builtins/zip.rs

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,15 @@
11
use super::PyTypeRef;
22
use crate::{
33
function::PosArgs,
4-
iterator,
4+
protocol::PyIter,
55
slots::{IteratorIterable, SlotConstructor, SlotIterator},
6-
PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, VirtualMachine,
6+
PyClassImpl, PyContext, PyRef, PyResult, PyValue, VirtualMachine,
77
};
88

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

1515
impl PyValue for PyZip {
@@ -19,13 +19,10 @@ impl PyValue for PyZip {
1919
}
2020

2121
impl SlotConstructor for PyZip {
22-
type Args = PosArgs;
22+
type Args = PosArgs<PyIter>;
2323

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

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

‎vm/src/frame.rs

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ use crate::{
1717
exceptions::{self, ExceptionCtor},
1818
function::FuncArgs,
1919
iterator,
20+
protocol::PyIter,
2021
scope::Scope,
2122
slots::PyComparisonOp,
2223
IdProtocol, ItemProtocol, PyMethod, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
@@ -857,8 +858,8 @@ impl ExecutingFrame<'_> {
857858
}
858859
bytecode::Instruction::GetIter => {
859860
let iterated_obj = self.pop_value();
860-
let iter_obj = iterator::get_iter(vm, iterated_obj)?;
861-
self.push_value(iter_obj);
861+
let iter_obj = iterated_obj.get_iter(vm)?;
862+
self.push_value(iter_obj.into_object());
862863
Ok(None)
863864
}
864865
bytecode::Instruction::GetAwaitable => {
@@ -1442,7 +1443,7 @@ impl ExecutingFrame<'_> {
14421443
fn _send(&self, coro: &PyObjectRef, val: PyObjectRef, vm: &VirtualMachine) -> PyResult {
14431444
match self.builtin_coro(coro) {
14441445
Some(coro) => coro.send(val, vm),
1445-
None if vm.is_none(&val) => iterator::call_next(vm, coro),
1446+
None if vm.is_none(&val) => PyIter::new(coro).next(vm),
14461447
None => {
14471448
let meth = vm.get_attribute(coro.clone(), "send")?;
14481449
vm.invoke(&meth, (val,))
@@ -1512,7 +1513,7 @@ impl ExecutingFrame<'_> {
15121513

15131514
/// The top of stack contains the iterator, lets push it forward
15141515
fn execute_for_iter(&mut self, vm: &VirtualMachine, target: bytecode::Label) -> FrameResult {
1515-
let top_of_stack = self.last_value();
1516+
let top_of_stack = PyIter::new(self.last_value());
15161517
let next_obj = iterator::get_next_object(vm, &top_of_stack);
15171518

15181519
// Check the next object:

‎vm/src/function/argument.rs

Lines changed: 3 additions & 3 deletions
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)