Skip to content

Commit 750687e

Browse files
youknowoneyouknowone
committed
PyIter protocol
1 parent 73aed46 commit 750687e

File tree

18 files changed

+272
-242
lines changed

18 files changed

+272
-242
lines changed

‎vm/src/builtins/dict.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -84,13 +84,13 @@ impl PyDict {
8484
return Err(vm.new_runtime_error("dict mutated during update".to_owned()));
8585
}
8686
} else if let Some(keys) = vm.get_method(dict_obj.clone(), "keys") {
87-
let keys = iterator::get_iter(vm, vm.invoke(&keys?, ())?)?;
87+
let keys = vm.invoke(&keys?, ())?.get_iter(vm)?;
8888
while let Some(key) = iterator::get_next_object(vm, &keys)? {
8989
let val = dict_obj.get_item(key.clone(), vm)?;
9090
dict.insert(vm, key, val)?;
9191
}
9292
} else {
93-
let iter = iterator::get_iter(vm, dict_obj)?;
93+
let iter = dict_obj.get_iter(vm)?;
9494
loop {
9595
fn err(vm: &VirtualMachine) -> PyBaseExceptionRef {
9696
vm.new_value_error("Iterator must have exactly two elements".to_owned())
@@ -99,7 +99,7 @@ impl PyDict {
9999
Some(obj) => obj,
100100
None => break,
101101
};
102-
let elem_iter = iterator::get_iter(vm, element)?;
102+
let elem_iter = element.get_iter(vm)?;
103103
let key = iterator::get_next_object(vm, &elem_iter)?.ok_or_else(|| err(vm))?;
104104
let value =
105105
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 {
@@ -31,19 +31,20 @@ impl PyValue for PyEnumerate {
3131
#[derive(FromArgs)]
3232
pub struct EnumerateArgs {
3333
#[pyarg(any)]
34-
iterable: PyObjectRef,
34+
iterator: PyIter,
3535
#[pyarg(any, optional)]
3636
start: OptionalArg<PyIntRef>,
3737
}
3838

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

42-
fn py_new(cls: PyTypeRef, args: Self::Args, vm: &VirtualMachine) -> PyResult {
43-
let counter = args
44-
.start
45-
.map_or_else(BigInt::zero, |start| start.as_bigint().clone());
46-
let iterator = iterator::get_iter(vm, args.iterable)?;
42+
fn py_new(
43+
cls: PyTypeRef,
44+
Self::Args { iterator, start }: Self::Args,
45+
vm: &VirtualMachine,
46+
) -> PyResult {
47+
let counter = start.map_or_else(BigInt::zero, |start| start.as_bigint().clone());
4748
PyEnumerate {
4849
counter: PyRwLock::new(counter),
4950
iterator,
@@ -58,7 +59,7 @@ impl PyEnumerate {}
5859
impl IteratorIterable for PyEnumerate {}
5960
impl SlotIterator for PyEnumerate {
6061
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
61-
let next_obj = iterator::call_next(vm, &zelf.iterator)?;
62+
let next_obj = zelf.iterator.next(vm)?;
6263
let mut counter = zelf.counter.write();
6364
let position = counter.clone();
6465
*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
function::{
2626
ArgCallable, ArgIterable, FuncArgs, KwArgs, OptionalArg, OptionalOption, PosArgs,
2727
},
28-
iterator, py_io,
28+
protocol::PyIter,
29+
py_io,
2930
readline::{Readline, ReadlineResult},
3031
scope::Scope,
3132
slots::PyComparisonOp,
@@ -416,14 +417,15 @@ mod decl {
416417
iter_target: PyObjectRef,
417418
sentinel: OptionalArg<PyObjectRef>,
418419
vm: &VirtualMachine,
419-
) -> PyResult {
420+
) -> PyResult<PyIter> {
420421
if let OptionalArg::Present(sentinel) = sentinel {
421422
let callable = ArgCallable::try_from_object(vm, iter_target)?;
422-
Ok(PyCallableIterator::new(callable, sentinel)
423+
let iterator = PyCallableIterator::new(callable, sentinel)
423424
.into_ref(vm)
424-
.into_object())
425+
.into_object();
426+
Ok(PyIter::new(iterator))
425427
} else {
426-
iterator::get_iter(vm, iter_target)
428+
iter_target.get_iter(vm)
427429
}
428430
}
429431

@@ -514,7 +516,7 @@ mod decl {
514516
default_value: OptionalArg<PyObjectRef>,
515517
vm: &VirtualMachine,
516518
) -> PyResult {
517-
iterator::call_next(vm, &iterator).or_else(|err| {
519+
PyIter::new(iterator).next(vm).or_else(|err| {
518520
if err.isinstance(&vm.ctx.exceptions.stop_iteration) {
519521
default_value.ok_or(err)
520522
} 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
@@ -16,6 +16,7 @@ use crate::{
1616
exceptions::{self, ExceptionCtor, PyBaseExceptionRef},
1717
function::FuncArgs,
1818
iterator,
19+
protocol::PyIter,
1920
scope::Scope,
2021
slots::PyComparisonOp,
2122
IdProtocol, ItemProtocol, PyMethod, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
@@ -856,8 +857,8 @@ impl ExecutingFrame<'_> {
856857
}
857858
bytecode::Instruction::GetIter => {
858859
let iterated_obj = self.pop_value();
859-
let iter_obj = iterator::get_iter(vm, iterated_obj)?;
860-
self.push_value(iter_obj);
860+
let iter_obj = iterated_obj.get_iter(vm)?;
861+
self.push_value(iter_obj.into_object());
861862
Ok(None)
862863
}
863864
bytecode::Instruction::GetAwaitable => {
@@ -1441,7 +1442,7 @@ impl ExecutingFrame<'_> {
14411442
fn _send(&self, coro: &PyObjectRef, val: PyObjectRef, vm: &VirtualMachine) -> PyResult {
14421443
match self.builtin_coro(coro) {
14431444
Some(coro) => coro.send(val, vm),
1444-
None if vm.is_none(&val) => iterator::call_next(vm, coro),
1445+
None if vm.is_none(&val) => PyIter::new(coro).next(vm),
14451446
None => {
14461447
let meth = vm.get_attribute(coro.clone(), "send")?;
14471448
vm.invoke(&meth, (val,))
@@ -1511,7 +1512,7 @@ impl ExecutingFrame<'_> {
15111512

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

15171518
// 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)