diff --git a/Lib/test/test_types.py b/Lib/test/test_types.py index 779ee6970..4ba8d5d33 100644 --- a/Lib/test/test_types.py +++ b/Lib/test/test_types.py @@ -1456,6 +1456,7 @@ class CoroutineTests(unittest.TestCase): self.assertIs(foo(), coro) self.assertIs(foo().__await__(), coro) + @unittest.skip("TODO: RUSTPYTHON, unittest.mock") def test_duck_gen(self): class GenLike: def send(self): pass diff --git a/derive/src/pyclass.rs b/derive/src/pyclass.rs index 5272778f0..e8aa209bf 100644 --- a/derive/src/pyclass.rs +++ b/derive/src/pyclass.rs @@ -92,6 +92,12 @@ pub(crate) fn impl_pyimpl( let mut context = ImplContext::default(); extract_items_into_context(&mut context, trai.items.iter_mut()); + let ExtractedImplAttrs { + with_impl, + with_slots, + .. + } = extract_impl_attrs(attr)?; + let getset_impl = &context.getset_items; let extend_impl = &context.impl_extend_items; let slots_impl = &context.extend_slots_items; @@ -104,11 +110,13 @@ pub(crate) fn impl_pyimpl( ) { #getset_impl #extend_impl + #with_impl #(#class_extensions)* } }, parse_quote! { fn __extend_slots(slots: &mut ::rustpython_vm::slots::PyTypeSlots) { + #with_slots #slots_impl } }, diff --git a/vm/src/builtins/asyncgenerator.rs b/vm/src/builtins/asyncgenerator.rs index 02717c52b..f2e679442 100644 --- a/vm/src/builtins/asyncgenerator.rs +++ b/vm/src/builtins/asyncgenerator.rs @@ -7,6 +7,7 @@ use crate::function::OptionalArg; use crate::pyobject::{ PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol, }; +use crate::slots::PyIter; use crate::vm::VirtualMachine; use crossbeam_utils::atomic::AtomicCell; @@ -172,22 +173,13 @@ impl PyValue for PyAsyncGenASend { } } -#[pyimpl] +#[pyimpl(with(PyIter))] impl PyAsyncGenASend { - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef, _vm: &VirtualMachine) -> PyRef { - zelf - } #[pymethod(name = "__await__")] fn r#await(zelf: PyRef, _vm: &VirtualMachine) -> PyRef { zelf } - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - self.send(vm.ctx.none(), vm) - } - #[pymethod] fn send(&self, val: PyObjectRef, vm: &VirtualMachine) -> PyResult { let val = match self.state.load() { @@ -253,6 +245,12 @@ impl PyAsyncGenASend { } } +impl PyIter for PyAsyncGenASend { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + zelf.send(vm.ctx.none(), vm) + } +} + #[pyclass(module = false, name = "async_generator_athrow")] #[derive(Debug)] pub(crate) struct PyAsyncGenAThrow { @@ -268,22 +266,13 @@ impl PyValue for PyAsyncGenAThrow { } } -#[pyimpl] +#[pyimpl(with(PyIter))] impl PyAsyncGenAThrow { - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef, _vm: &VirtualMachine) -> PyRef { - zelf - } #[pymethod(name = "__await__")] fn r#await(zelf: PyRef, _vm: &VirtualMachine) -> PyRef { zelf } - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - self.send(vm.ctx.none(), vm) - } - #[pymethod] fn send(&self, val: PyObjectRef, vm: &VirtualMachine) -> PyResult { match self.state.load() { @@ -397,6 +386,12 @@ impl PyAsyncGenAThrow { } } +impl PyIter for PyAsyncGenAThrow { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + zelf.send(vm.ctx.none(), vm) + } +} + pub fn init(ctx: &PyContext) { PyAsyncGen::extend_class(ctx, &ctx.types.async_generator); PyAsyncGenASend::extend_class(ctx, &ctx.types.async_generator_asend); diff --git a/vm/src/builtins/bytearray.rs b/vm/src/builtins/bytearray.rs index 284251e65..7f399a381 100644 --- a/vm/src/builtins/bytearray.rs +++ b/vm/src/builtins/bytearray.rs @@ -21,7 +21,9 @@ use crate::pyobject::{ PyIterable, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol, }; use crate::sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex}; -use crate::slots::{BufferProtocol, Comparable, Hashable, PyComparisonOp, Unhashable}; +use crate::slots::{ + BufferProtocol, Comparable, Hashable, Iterable, PyComparisonOp, PyIter, Unhashable, +}; use crate::vm::VirtualMachine; use bstr::ByteSlice; use crossbeam_utils::atomic::AtomicCell; @@ -97,7 +99,7 @@ pub(crate) fn init(context: &PyContext) { PyByteArrayIterator::extend_class(context, &context.types.bytearray_iterator_type); } -#[pyimpl(flags(BASETYPE), with(Hashable, Comparable, BufferProtocol))] +#[pyimpl(flags(BASETYPE), with(Hashable, Comparable, BufferProtocol, Iterable))] impl PyByteArray { #[pyslot] fn tp_new( @@ -123,14 +125,6 @@ impl PyByteArray { size_of::() + self.borrow_value().len() * size_of::() } - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyByteArrayIterator { - PyByteArrayIterator { - position: AtomicCell::new(0), - bytearray: zelf, - } - } - #[pymethod(name = "__add__")] fn add(&self, other: PyBytesLike, vm: &VirtualMachine) -> PyObjectRef { vm.ctx @@ -718,6 +712,16 @@ impl<'a> ResizeGuard<'a> for PyByteArray { impl Unhashable for PyByteArray {} +impl Iterable for PyByteArray { + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + Ok(PyByteArrayIterator { + position: AtomicCell::new(0), + bytearray: zelf, + } + .into_object(vm)) + } +} + // fn set_value(obj: &PyObjectRef, value: Vec) { // obj.borrow_mut().kind = PyObjectPayload::Bytes { value }; // } @@ -735,20 +739,15 @@ impl PyValue for PyByteArrayIterator { } } -#[pyimpl] -impl PyByteArrayIterator { - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let pos = self.position.fetch_add(1); - if let Some(&ret) = self.bytearray.borrow_value().elements.get(pos) { - Ok(ret) +#[pyimpl(with(PyIter))] +impl PyByteArrayIterator {} +impl PyIter for PyByteArrayIterator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let pos = zelf.position.fetch_add(1); + if let Some(&ret) = zelf.bytearray.borrow_value().elements.get(pos) { + Ok(ret.into_pyobject(vm)) } else { Err(vm.new_stop_iteration()) } } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } diff --git a/vm/src/builtins/bytes.rs b/vm/src/builtins/bytes.rs index caf21fb72..9109a2d7d 100644 --- a/vm/src/builtins/bytes.rs +++ b/vm/src/builtins/bytes.rs @@ -19,7 +19,7 @@ use crate::pyobject::{ BorrowValue, Either, IntoPyObject, PyClassImpl, PyComparisonValue, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol, }; -use crate::slots::{BufferProtocol, Comparable, Hashable, PyComparisonOp}; +use crate::slots::{BufferProtocol, Comparable, Hashable, Iterable, PyComparisonOp, PyIter}; use crate::vm::VirtualMachine; use crate::{ anystr::{self, AnyStr}, @@ -96,7 +96,7 @@ pub(crate) fn init(context: &PyContext) { PyBytesIterator::extend_class(context, &context.types.bytes_iterator_type); } -#[pyimpl(flags(BASETYPE), with(Hashable, Comparable, BufferProtocol))] +#[pyimpl(flags(BASETYPE), with(Hashable, Comparable, BufferProtocol, Iterable))] impl PyBytes { #[pyslot] fn tp_new( @@ -117,14 +117,6 @@ impl PyBytes { self.inner.len() } - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyBytesIterator { - PyBytesIterator { - position: AtomicCell::new(0), - bytes: zelf, - } - } - #[pymethod(name = "__sizeof__")] fn sizeof(&self) -> PyResult { Ok(size_of::() + self.inner.elements.len() * size_of::()) @@ -545,6 +537,16 @@ impl Comparable for PyBytes { } } +impl Iterable for PyBytes { + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + Ok(PyBytesIterator { + position: AtomicCell::new(0), + bytes: zelf, + } + .into_object(vm)) + } +} + #[pyclass(module = false, name = "bytes_iterator")] #[derive(Debug)] pub struct PyBytesIterator { @@ -558,22 +560,17 @@ impl PyValue for PyBytesIterator { } } -#[pyimpl] -impl PyBytesIterator { - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let pos = self.position.fetch_add(1); - if let Some(&ret) = self.bytes.borrow_value().get(pos) { - Ok(ret) +#[pyimpl(with(PyIter))] +impl PyBytesIterator {} +impl PyIter for PyBytesIterator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let pos = zelf.position.fetch_add(1); + if let Some(&ret) = zelf.bytes.borrow_value().get(pos) { + Ok(vm.ctx.new_int(ret)) } else { Err(vm.new_stop_iteration()) } } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } impl TryFromObject for PyBytes { diff --git a/vm/src/builtins/coroutine.rs b/vm/src/builtins/coroutine.rs index d7157b63c..cba341f18 100644 --- a/vm/src/builtins/coroutine.rs +++ b/vm/src/builtins/coroutine.rs @@ -5,6 +5,7 @@ use crate::coroutine::{Coro, Variant}; use crate::frame::FrameRef; use crate::function::OptionalArg; use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; +use crate::slots::PyIter; use crate::vm::VirtualMachine; type PyCoroutineRef = PyRef; @@ -21,7 +22,7 @@ impl PyValue for PyCoroutine { } } -#[pyimpl] +#[pyimpl(with(PyIter))] impl PyCoroutine { pub fn as_coro(&self) -> &Coro { &self.inner @@ -95,6 +96,12 @@ impl PyCoroutine { } } +impl PyIter for PyCoroutine { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + zelf.send(vm.ctx.none(), vm) + } +} + #[pyclass(module = false, name = "coroutine_wrapper")] #[derive(Debug)] pub struct PyCoroutineWrapper { @@ -109,16 +116,6 @@ impl PyValue for PyCoroutineWrapper { #[pyimpl] impl PyCoroutineWrapper { - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - self.coro.send(vm.ctx.none(), vm) - } - #[pymethod] fn send(&self, val: PyObjectRef, vm: &VirtualMachine) -> PyResult { self.coro.send(val, vm) diff --git a/vm/src/builtins/dict.rs b/vm/src/builtins/dict.rs index fd0968c6c..d716a2fd0 100644 --- a/vm/src/builtins/dict.rs +++ b/vm/src/builtins/dict.rs @@ -14,7 +14,7 @@ use crate::pyobject::{ PyClassImpl, PyComparisonValue, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol, }; -use crate::slots::{Comparable, Hashable, PyComparisonOp, Unhashable}; +use crate::slots::{Comparable, Hashable, Iterable, PyComparisonOp, PyIter, Unhashable}; use crate::vm::{ReprGuard, VirtualMachine}; pub type DictContentType = dictdatatype::Dict; @@ -50,7 +50,7 @@ impl PyValue for PyDict { // Python dict methods: #[allow(clippy::len_without_is_empty)] -#[pyimpl(with(Hashable, Comparable), flags(BASETYPE))] +#[pyimpl(with(Hashable, Comparable, Iterable), flags(BASETYPE))] impl PyDict { #[pyslot] fn tp_new(class: PyTypeRef, _args: FuncArgs, vm: &VirtualMachine) -> PyResult> { @@ -83,13 +83,13 @@ impl PyDict { dict.insert(vm, key, value)?; } } else if let Some(keys) = vm.get_method(dict_obj.clone(), "keys") { - let keys = iterator::get_iter(vm, &vm.invoke(&keys?, ())?)?; + let keys = iterator::get_iter(vm, vm.invoke(&keys?, ())?)?; while let Some(key) = iterator::get_next_object(vm, &keys)? { let val = dict_obj.get_item(key.clone(), vm)?; dict.insert(vm, key, val)?; } } else { - let iter = iterator::get_iter(vm, &dict_obj)?; + let iter = iterator::get_iter(vm, dict_obj)?; loop { fn err(vm: &VirtualMachine) -> PyBaseExceptionRef { vm.new_value_error("Iterator must have exactly two elements".to_owned()) @@ -98,7 +98,7 @@ impl PyDict { Some(obj) => obj, None => break, }; - let elem_iter = iterator::get_iter(vm, &element)?; + let elem_iter = iterator::get_iter(vm, element)?; let key = iterator::get_next_object(vm, &elem_iter)?.ok_or_else(|| err(vm))?; let value = iterator::get_next_object(vm, &elem_iter)?.ok_or_else(|| err(vm))?; @@ -234,11 +234,6 @@ impl PyDict { self.entries.clear() } - #[pymethod(magic)] - fn iter(zelf: PyRef) -> PyDictKeyIterator { - PyDictKeyIterator::new(zelf) - } - #[pymethod] fn keys(zelf: PyRef) -> PyDictKeys { PyDictKeys::new(zelf) @@ -420,6 +415,12 @@ impl Comparable for PyDict { impl Unhashable for PyDict {} +impl Iterable for PyDict { + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + Ok(PyDictKeyIterator::new(zelf).into_object(vm)) + } +} + impl PyDictRef { /// Return an optional inner item, or an error (can be key error as well) #[inline] @@ -579,17 +580,12 @@ macro_rules! dict_iterator { pub dict: PyDictRef, } - #[pyimpl(with(Comparable))] + #[pyimpl(with(Comparable, Iterable))] impl $name { fn new(dict: PyDictRef) -> Self { $name { dict } } - #[pymethod(name = "__iter__")] - fn iter(&self) -> $iter_name { - $iter_name::new(self.dict.clone()) - } - #[pymethod(name = "__len__")] fn len(&self) -> usize { self.dict.clone().len() @@ -616,6 +612,12 @@ macro_rules! dict_iterator { } } + impl Iterable for $name { + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + Ok($iter_name::new(zelf.dict.clone()).into_object(vm)) + } + } + impl Comparable for $name { fn cmp( zelf: &PyRef, @@ -658,7 +660,13 @@ macro_rules! dict_iterator { pub position: AtomicCell, } - #[pyimpl] + impl PyValue for $iter_name { + fn class(vm: &VirtualMachine) -> &PyTypeRef { + &vm.ctx.types.$iter_class + } + } + + #[pyimpl(with(PyIter))] impl $iter_name { fn new(dict: PyDictRef) -> Self { $iter_name { @@ -668,38 +676,28 @@ macro_rules! dict_iterator { } } - #[pymethod(name = "__next__")] - #[allow(clippy::redundant_closure_call)] - fn next(&self, vm: &VirtualMachine) -> PyResult { - if self.dict.entries.has_changed_size(&self.size) { - return Err( - vm.new_runtime_error("dictionary changed size during iteration".to_owned()) - ); - } - let mut position = self.position.load(); - match self.dict.entries.next_entry(&mut position) { - Some((key, value)) => { - self.position.store(position); - Ok($result_fn(vm, key, value)) - } - None => Err(vm.new_stop_iteration()), - } - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - #[pymethod(name = "__length_hint__")] fn length_hint(&self) -> usize { self.dict.entries.len_from_entry_index(self.position.load()) } } - impl PyValue for $iter_name { - fn class(vm: &VirtualMachine) -> &PyTypeRef { - &vm.ctx.types.$iter_class + impl PyIter for $iter_name { + #[allow(clippy::redundant_closure_call)] + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + if zelf.dict.entries.has_changed_size(&zelf.size) { + return Err( + vm.new_runtime_error("dictionary changed size during iteration".to_owned()) + ); + } + let mut position = zelf.position.load(); + match zelf.dict.entries.next_entry(&mut position) { + Some((key, value)) => { + zelf.position.store(position); + Ok($result_fn(vm, key, value)) + } + None => Err(vm.new_stop_iteration()), + } } } @@ -711,7 +709,13 @@ macro_rules! dict_iterator { pub position: AtomicCell, } - #[pyimpl] + impl PyValue for $reverse_iter_name { + fn class(vm: &VirtualMachine) -> &PyTypeRef { + &vm.ctx.types.$reverse_iter_class + } + } + + #[pyimpl(with(PyIter))] impl $reverse_iter_name { fn new(dict: PyDictRef) -> Self { $reverse_iter_name { @@ -721,41 +725,31 @@ macro_rules! dict_iterator { } } - #[pymethod(name = "__next__")] - #[allow(clippy::redundant_closure_call)] - fn next(&self, vm: &VirtualMachine) -> PyResult { - if self.dict.entries.has_changed_size(&self.size) { - return Err( - vm.new_runtime_error("dictionary changed size during iteration".to_owned()) - ); - } - let count = self.position.fetch_add(1); - match self.dict.len().checked_sub(count) { - Some(mut pos) => { - let (key, value) = self.dict.entries.next_entry(&mut pos).unwrap(); - Ok($result_fn(vm, key, value)) - } - None => { - self.position.store(std::isize::MAX as usize); - Err(vm.new_stop_iteration()) - } - } - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - #[pymethod(name = "__length_hint__")] fn length_hint(&self) -> usize { self.dict.entries.len_from_entry_index(self.position.load()) } } - impl PyValue for $reverse_iter_name { - fn class(vm: &VirtualMachine) -> &PyTypeRef { - &vm.ctx.types.$reverse_iter_class + impl PyIter for $reverse_iter_name { + #[allow(clippy::redundant_closure_call)] + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + if zelf.dict.entries.has_changed_size(&zelf.size) { + return Err( + vm.new_runtime_error("dictionary changed size during iteration".to_owned()) + ); + } + let count = zelf.position.fetch_add(1); + match zelf.dict.len().checked_sub(count) { + Some(mut pos) => { + let (key, value) = zelf.dict.entries.next_entry(&mut pos).unwrap(); + Ok($result_fn(vm, key, value)) + } + None => { + zelf.position.store(std::isize::MAX as usize); + Err(vm.new_stop_iteration()) + } + } } } }; diff --git a/vm/src/builtins/enumerate.rs b/vm/src/builtins/enumerate.rs index c11342c05..b0e0ee7de 100644 --- a/vm/src/builtins/enumerate.rs +++ b/vm/src/builtins/enumerate.rs @@ -1,5 +1,4 @@ use crate::common::lock::PyRwLock; -use std::ops::AddAssign; use num_bigint::BigInt; use num_traits::Zero; @@ -8,7 +7,10 @@ use super::int::PyIntRef; use super::pytype::PyTypeRef; use crate::function::OptionalArg; use crate::iterator; -use crate::pyobject::{BorrowValue, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; +use crate::pyobject::{ + BorrowValue, IntoPyObject, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, +}; +use crate::slots::PyIter; use crate::vm::VirtualMachine; #[pyclass(module = false, name = "enumerate")] @@ -24,7 +26,7 @@ impl PyValue for PyEnumerate { } } -#[pyimpl] +#[pyimpl(with(PyIter))] impl PyEnumerate { #[pyslot] fn tp_new( @@ -38,26 +40,22 @@ impl PyEnumerate { OptionalArg::Missing => BigInt::zero(), }; - let iterator = iterator::get_iter(vm, &iterable)?; + let iterator = iterator::get_iter(vm, iterable)?; PyEnumerate { counter: PyRwLock::new(counter), iterator, } .into_ref_with_type(vm, cls) } +} - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult<(BigInt, PyObjectRef)> { - let next_obj = iterator::call_next(vm, &self.iterator)?; - let mut counter = self.counter.write(); +impl PyIter for PyEnumerate { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let next_obj = iterator::call_next(vm, &zelf.iterator)?; + let mut counter = zelf.counter.write(); let position = counter.clone(); - AddAssign::add_assign(&mut counter as &mut BigInt, 1); - Ok((position, next_obj)) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf + *counter += 1; + Ok((position, next_obj).into_pyobject(vm)) } } diff --git a/vm/src/builtins/filter.rs b/vm/src/builtins/filter.rs index 067d41b8d..d0460c2a8 100644 --- a/vm/src/builtins/filter.rs +++ b/vm/src/builtins/filter.rs @@ -2,6 +2,7 @@ use super::pybool; use super::pytype::PyTypeRef; use crate::iterator; use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; +use crate::slots::PyIter; use crate::vm::VirtualMachine; /// filter(function or None, iterable) --> filter object @@ -21,7 +22,7 @@ impl PyValue for PyFilter { } } -#[pyimpl(flags(BASETYPE))] +#[pyimpl(with(PyIter), flags(BASETYPE))] impl PyFilter { #[pyslot] fn tp_new( @@ -30,7 +31,7 @@ impl PyFilter { iterable: PyObjectRef, vm: &VirtualMachine, ) -> PyResult> { - let iterator = iterator::get_iter(vm, &iterable)?; + let iterator = iterator::get_iter(vm, iterable)?; Self { predicate: function, @@ -38,11 +39,12 @@ impl PyFilter { } .into_ref_with_type(vm, cls) } +} - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let predicate = &self.predicate; - let iterator = &self.iterator; +impl PyIter for PyFilter { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let predicate = &zelf.predicate; + let iterator = &zelf.iterator; loop { let next_obj = iterator::call_next(vm, iterator)?; let predicate_value = if vm.is_none(predicate) { @@ -57,11 +59,6 @@ impl PyFilter { } } } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } pub fn init(context: &PyContext) { diff --git a/vm/src/builtins/generator.rs b/vm/src/builtins/generator.rs index b8213b006..d33105772 100644 --- a/vm/src/builtins/generator.rs +++ b/vm/src/builtins/generator.rs @@ -8,6 +8,7 @@ use crate::coroutine::{Coro, Variant}; use crate::frame::FrameRef; use crate::function::OptionalArg; use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; +use crate::slots::PyIter; use crate::vm::VirtualMachine; #[pyclass(module = false, name = "generator")] @@ -22,7 +23,7 @@ impl PyValue for PyGenerator { } } -#[pyimpl] +#[pyimpl(with(PyIter))] impl PyGenerator { pub fn as_coro(&self) -> &Coro { &self.inner @@ -41,16 +42,6 @@ impl PyGenerator { vm.ctx.none() } - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - self.send(vm.ctx.none(), vm) - } - #[pymethod] fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult { self.inner.send(value, vm) @@ -95,6 +86,12 @@ impl PyGenerator { } } +impl PyIter for PyGenerator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + zelf.send(vm.ctx.none(), vm) + } +} + pub fn init(ctx: &PyContext) { PyGenerator::extend_class(ctx, &ctx.types.generator_type); } diff --git a/vm/src/builtins/iter.rs b/vm/src/builtins/iter.rs index 1081b16f3..6b77321a6 100644 --- a/vm/src/builtins/iter.rs +++ b/vm/src/builtins/iter.rs @@ -8,6 +8,7 @@ use super::pytype::PyTypeRef; use crate::pyobject::{ PyCallable, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol, }; +use crate::slots::PyIter; use crate::vm::VirtualMachine; #[pyclass(module = false, name = "iter")] @@ -24,7 +25,7 @@ impl PyValue for PySequenceIterator { } } -#[pyimpl] +#[pyimpl(with(PyIter))] impl PySequenceIterator { pub fn new_forward(obj: PyObjectRef) -> Self { Self { @@ -42,28 +43,6 @@ impl PySequenceIterator { } } - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let step: isize = if self.reversed { -1 } else { 1 }; - let pos = self.position.fetch_add(step); - if pos >= 0 { - match vm.call_method(&self.obj, "__getitem__", (pos,)) { - Err(ref e) if e.isinstance(&vm.ctx.exceptions.index_error) => { - Err(vm.new_stop_iteration()) - } - // also catches stop_iteration => stop_iteration - ret => ret, - } - } else { - Err(vm.new_stop_iteration()) - } - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - #[pymethod(name = "__length_hint__")] fn length_hint(&self, vm: &VirtualMachine) -> PyResult { let pos = self.position.load(); @@ -77,6 +56,24 @@ impl PySequenceIterator { } } +impl PyIter for PySequenceIterator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let step: isize = if zelf.reversed { -1 } else { 1 }; + let pos = zelf.position.fetch_add(step); + if pos >= 0 { + match vm.call_method(&zelf.obj, "__getitem__", (pos,)) { + Err(ref e) if e.isinstance(&vm.ctx.exceptions.index_error) => { + Err(vm.new_stop_iteration()) + } + // also catches stop_iteration => stop_iteration + ret => ret, + } + } else { + Err(vm.new_stop_iteration()) + } + } +} + #[pyclass(module = false, name = "callable_iterator")] #[derive(Debug)] pub struct PyCallableIterator { @@ -91,7 +88,7 @@ impl PyValue for PyCallableIterator { } } -#[pyimpl] +#[pyimpl(with(PyIter))] impl PyCallableIterator { pub fn new(callable: PyCallable, sentinel: PyObjectRef) -> Self { Self { @@ -100,27 +97,23 @@ impl PyCallableIterator { done: AtomicCell::new(false), } } +} - #[pymethod(magic)] - fn next(&self, vm: &VirtualMachine) -> PyResult { - if self.done.load() { +impl PyIter for PyCallableIterator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + if zelf.done.load() { return Err(vm.new_stop_iteration()); } - let ret = self.callable.invoke((), vm)?; + let ret = zelf.callable.invoke((), vm)?; - if vm.bool_eq(&ret, &self.sentinel)? { - self.done.store(true); + if vm.bool_eq(&ret, &zelf.sentinel)? { + zelf.done.store(true); Err(vm.new_stop_iteration()) } else { Ok(ret) } } - - #[pymethod(magic)] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } pub fn init(context: &PyContext) { diff --git a/vm/src/builtins/list.rs b/vm/src/builtins/list.rs index ff9a7f975..c77228de5 100644 --- a/vm/src/builtins/list.rs +++ b/vm/src/builtins/list.rs @@ -15,7 +15,7 @@ use crate::pyobject::{ }; use crate::sequence::{self, SimpleSeq}; use crate::sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex}; -use crate::slots::{Comparable, Hashable, PyComparisonOp, Unhashable}; +use crate::slots::{Comparable, Hashable, Iterable, PyComparisonOp, PyIter, Unhashable}; use crate::vm::{ReprGuard, VirtualMachine}; /// Built-in mutable sequence. @@ -79,7 +79,7 @@ pub(crate) struct SortOptions { pub type PyListRef = PyRef; -#[pyimpl(with(Hashable, Comparable), flags(BASETYPE))] +#[pyimpl(with(Iterable, Hashable, Comparable), flags(BASETYPE))] impl PyList { #[pymethod] pub(crate) fn append(&self, x: PyObjectRef) { @@ -175,14 +175,6 @@ impl PyList { Ok(result) } - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyListIterator { - PyListIterator { - position: AtomicCell::new(0), - list: zelf, - } - } - #[pymethod(name = "__setitem__")] fn setitem( &self, @@ -383,6 +375,16 @@ impl PyList { } } +impl Iterable for PyList { + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + Ok(PyListIterator { + position: AtomicCell::new(0), + list: zelf, + } + .into_object(vm)) + } +} + impl Comparable for PyList { fn cmp( zelf: &PyRef, @@ -442,24 +444,8 @@ impl PyValue for PyListIterator { } } -#[pyimpl] +#[pyimpl(with(PyIter))] impl PyListIterator { - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let list = self.list.borrow_value(); - let pos = self.position.fetch_add(1); - if let Some(obj) = list.get(pos) { - Ok(obj.clone()) - } else { - Err(vm.new_stop_iteration()) - } - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - #[pymethod(name = "__length_hint__")] fn length_hint(&self) -> usize { let list = self.list.borrow_value(); @@ -468,6 +454,18 @@ impl PyListIterator { } } +impl PyIter for PyListIterator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let list = zelf.list.borrow_value(); + let pos = zelf.position.fetch_add(1); + if let Some(obj) = list.get(pos) { + Ok(obj.clone()) + } else { + Err(vm.new_stop_iteration()) + } + } +} + #[pyclass(module = false, name = "list_reverseiterator")] #[derive(Debug)] pub struct PyListReverseIterator { @@ -481,12 +479,18 @@ impl PyValue for PyListReverseIterator { } } -#[pyimpl] +#[pyimpl(with(PyIter))] impl PyListReverseIterator { - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let list = self.list.borrow_value(); - let pos = self.position.fetch_sub(1); + #[pymethod(name = "__length_hint__")] + fn length_hint(&self) -> usize { + std::cmp::max(self.position.load(), 0) as usize + } +} + +impl PyIter for PyListReverseIterator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let list = zelf.list.borrow_value(); + let pos = zelf.position.fetch_sub(1); if pos > 0 { if let Some(ret) = list.get(pos as usize - 1) { return Ok(ret.clone()); @@ -494,16 +498,6 @@ impl PyListReverseIterator { } Err(vm.new_stop_iteration()) } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - - #[pymethod(name = "__length_hint__")] - fn length_hint(&self) -> usize { - std::cmp::max(self.position.load(), 0) as usize - } } pub fn init(context: &PyContext) { diff --git a/vm/src/builtins/make_module.rs b/vm/src/builtins/make_module.rs index e093f076d..f45a87c45 100644 --- a/vm/src/builtins/make_module.rs +++ b/vm/src/builtins/make_module.rs @@ -434,7 +434,7 @@ mod decl { .into_ref(vm) .into_object()) } else { - iterator::get_iter(vm, &iter_target) + iterator::get_iter(vm, iter_target) } } @@ -526,19 +526,16 @@ mod decl { default_value: OptionalArg, vm: &VirtualMachine, ) -> PyResult { - match vm.call_method(&iterator, "__next__", ()) { - Ok(value) => Ok(value), - Err(value) => { - if value.isinstance(&vm.ctx.exceptions.stop_iteration) { - match default_value { - OptionalArg::Missing => Err(value), - OptionalArg::Present(value) => Ok(value), - } - } else { - Err(value) + iterator::call_next(vm, &iterator).or_else(|err| { + if err.isinstance(&vm.ctx.exceptions.stop_iteration) { + match default_value { + OptionalArg::Missing => Err(err), + OptionalArg::Present(value) => Ok(value), } + } else { + Err(err) } - } + }) } #[pyfunction] diff --git a/vm/src/builtins/map.rs b/vm/src/builtins/map.rs index ae21dc8c5..21fcf237e 100644 --- a/vm/src/builtins/map.rs +++ b/vm/src/builtins/map.rs @@ -2,6 +2,7 @@ use super::pytype::PyTypeRef; use crate::function::Args; use crate::iterator; use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; +use crate::slots::PyIter; use crate::vm::VirtualMachine; /// map(func, *iterables) --> map object @@ -21,7 +22,7 @@ impl PyValue for PyMap { } } -#[pyimpl(flags(BASETYPE))] +#[pyimpl(with(PyIter), flags(BASETYPE))] impl PyMap { #[pyslot] fn tp_new( @@ -32,7 +33,7 @@ impl PyMap { ) -> PyResult> { let iterators = iterables .into_iter() - .map(|iterable| iterator::get_iter(vm, &iterable)) + .map(|iterable| iterator::get_iter(vm, iterable)) .collect::, _>>()?; PyMap { mapper: function, @@ -41,23 +42,6 @@ impl PyMap { .into_ref_with_type(vm, cls) } - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let next_objs = self - .iterators - .iter() - .map(|iterator| iterator::call_next(vm, iterator)) - .collect::, _>>()?; - - // the mapper itself can raise StopIteration which does stop the map iteration - vm.invoke(&self.mapper, next_objs) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - #[pymethod(name = "__length_hint__")] fn length_hint(&self, vm: &VirtualMachine) -> PyResult { self.iterators.iter().try_fold(0, |prev, cur| { @@ -68,6 +52,19 @@ impl PyMap { } } +impl PyIter for PyMap { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let next_objs = zelf + .iterators + .iter() + .map(|iterator| iterator::call_next(vm, iterator)) + .collect::, _>>()?; + + // the mapper itself can raise StopIteration which does stop the map iteration + vm.invoke(&zelf.mapper, next_objs) + } +} + pub fn init(context: &PyContext) { PyMap::extend_class(context, &context.types.map_type); } diff --git a/vm/src/builtins/mappingproxy.rs b/vm/src/builtins/mappingproxy.rs index 21ce7c776..df88ee2a8 100644 --- a/vm/src/builtins/mappingproxy.rs +++ b/vm/src/builtins/mappingproxy.rs @@ -7,6 +7,7 @@ use crate::pyobject::{ BorrowValue, IntoPyObject, ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, }; +use crate::slots::Iterable; use crate::vm::VirtualMachine; #[pyclass(module = false, name = "mappingproxy")] @@ -35,7 +36,7 @@ impl PyMappingProxy { } } -#[pyimpl] +#[pyimpl(with(Iterable))] impl PyMappingProxy { #[pyslot] fn tp_new(cls: PyTypeRef, mapping: PyObjectRef, vm: &VirtualMachine) -> PyResult> { @@ -85,17 +86,6 @@ impl PyMappingProxy { } } - #[pymethod(name = "__iter__")] - pub fn iter(&self, vm: &VirtualMachine) -> PyResult { - let obj = match &self.mapping { - MappingProxyInner::Dict(d) => d.clone(), - MappingProxyInner::Class(c) => { - // TODO: something that's much more efficient than this - PyDict::from_attributes(c.attributes.read().clone(), vm)?.into_pyobject(vm) - } - }; - iterator::get_iter(vm, &obj) - } #[pymethod] pub fn items(&self, vm: &VirtualMachine) -> PyResult { let obj = match &self.mapping { @@ -136,6 +126,18 @@ impl PyMappingProxy { } } } +impl Iterable for PyMappingProxy { + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + let obj = match &zelf.mapping { + MappingProxyInner::Dict(d) => d.clone(), + MappingProxyInner::Class(c) => { + // TODO: something that's much more efficient than this + PyDict::from_attributes(c.attributes.read().clone(), vm)?.into_pyobject(vm) + } + }; + iterator::get_iter(vm, obj) + } +} pub fn init(context: &PyContext) { PyMappingProxy::extend_class(context, &context.types.mappingproxy_type) diff --git a/vm/src/builtins/pystr.rs b/vm/src/builtins/pystr.rs index 2adf4629a..45ebe0100 100644 --- a/vm/src/builtins/pystr.rs +++ b/vm/src/builtins/pystr.rs @@ -26,7 +26,7 @@ use crate::pyobject::{ TypeProtocol, }; use crate::sliceable::PySliceableSequence; -use crate::slots::{Comparable, Hashable, PyComparisonOp}; +use crate::slots::{Comparable, Hashable, Iterable, PyComparisonOp, PyIter}; use crate::VirtualMachine; use rustpython_common::hash; @@ -119,12 +119,13 @@ impl PyValue for PyStrIterator { } } -#[pyimpl] -impl PyStrIterator { - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let value = &*self.string.value; - let start = self.position.load(); +#[pyimpl(with(PyIter))] +impl PyStrIterator {} + +impl PyIter for PyStrIterator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let value = &*zelf.string.value; + let start = zelf.position.load(); if start == value.len() { return Err(vm.new_stop_iteration()); } @@ -139,18 +140,13 @@ impl PyStrIterator { end.unwrap_or(start + 4) }; - if let Ok(_stored) = self.position.compare_exchange(start, end) { - Ok(value[start..end].to_owned()) + if zelf.position.compare_exchange(start, end).is_ok() { + Ok(value[start..end].into_pyobject(vm)) } else { // already taken from elsewhere - self.next(vm) + Self::next(zelf, vm) } } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } #[pyclass(module = false, name = "str_reverseiterator")] @@ -166,12 +162,13 @@ impl PyValue for PyStrReverseIterator { } } -#[pyimpl] -impl PyStrReverseIterator { - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let value = &*self.string.value; - let end = self.position.load(); +#[pyimpl(with(PyIter))] +impl PyStrReverseIterator {} + +impl PyIter for PyStrReverseIterator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let value = &*zelf.string.value; + let end = zelf.position.load(); if end == 0 { return Err(vm.new_stop_iteration()); } @@ -186,18 +183,13 @@ impl PyStrReverseIterator { start.unwrap_or(end - 4) }; - let stored = self.position.swap(start); + let stored = zelf.position.swap(start); if end != stored { // already taken from elsewhere - return self.next(vm); + return Self::next(zelf, vm); } - Ok(value[start..end].to_owned()) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf + Ok(value[start..end].into_pyobject(vm)) } } @@ -211,7 +203,7 @@ struct StrArgs { errors: OptionalArg, } -#[pyimpl(flags(BASETYPE), with(Hashable, Comparable))] +#[pyimpl(flags(BASETYPE), with(Hashable, Comparable, Iterable))] impl PyStr { #[pyslot] fn tp_new(cls: PyTypeRef, args: StrArgs, vm: &VirtualMachine) -> PyResult { @@ -1049,14 +1041,6 @@ impl PyStr { encode_string(zelf, args.encoding, args.errors, vm) } - #[pymethod(magic)] - fn iter(zelf: PyRef) -> PyStrIterator { - PyStrIterator { - position: AtomicCell::new(0), - string: zelf, - } - } - #[pymethod(magic)] fn reversed(zelf: PyRef) -> PyStrReverseIterator { let end = zelf.value.len(); @@ -1091,6 +1075,16 @@ impl Comparable for PyStr { } } +impl Iterable for PyStr { + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + Ok(PyStrIterator { + position: AtomicCell::new(0), + string: zelf, + } + .into_object(vm)) + } +} + #[derive(FromArgs)] struct EncodeArgs { #[pyarg(any, default)] diff --git a/vm/src/builtins/pytype.rs b/vm/src/builtins/pytype.rs index 49f8e3c0e..cdbe2aec5 100644 --- a/vm/src/builtins/pytype.rs +++ b/vm/src/builtins/pytype.rs @@ -143,8 +143,13 @@ impl PyType { attributes } - pub(crate) fn update_slot(&self, name: &str) { + pub(crate) fn update_slot(&self, name: &str, add: bool) { // self is the resolved class in get_class_magic + macro_rules! update_slot { + ($name:ident, $func:expr) => {{ + self.slots.$name.store(if add { Some($func) } else { None }); + }}; + } match name { "__call__" => { let func: slots::GenericMethod = |zelf, args, vm| { @@ -152,14 +157,14 @@ impl PyType { let magic = vm.call_if_get_descriptor(magic, zelf.clone())?; vm.invoke(&magic, args) } as _; - self.slots.call.store(Some(func)) + update_slot!(call, func); } "__get__" => { let func: slots::DescrGetFunc = |zelf, obj, cls, vm| { let magic = get_class_magic(&zelf, "__get__"); vm.invoke(&magic, (zelf, obj, cls)) } as _; - self.slots.descr_get.store(Some(func)) + update_slot!(descr_get, func); } "__hash__" => { let func: slots::HashFunc = |zelf, vm| { @@ -173,7 +178,7 @@ impl PyType { .new_type_error("__hash__ method should return an integer".to_owned())), } } as _; - self.slots.hash.store(Some(func)); + update_slot!(hash, func); } "__del__" => { let func: slots::DelFunc = |zelf, vm| { @@ -181,22 +186,37 @@ impl PyType { let _ = vm.invoke(&magic, vec![zelf.clone()])?; Ok(()) } as _; - self.slots.del.store(Some(func)); + update_slot!(del, func); } "__eq__" | "__ne__" | "__le__" | "__lt__" | "__ge__" | "__gt__" => { let func: slots::CmpFunc = |zelf, other, op, vm| { + // TODO: differentiate between op types; get_class_magic would panic let magic = get_class_magic(&zelf, op.method_name()); vm.invoke(&magic, vec![zelf.clone(), other.clone()]) .map(Either::A) } as _; - self.slots.cmp.store(Some(func)) + update_slot!(cmp, func); } "__getattribute__" => { let func: slots::GetattroFunc = |zelf, name, vm| { let magic = get_class_magic(&zelf, "__getattribute__"); vm.invoke(&magic, (zelf, name)) }; - self.slots.getattro.store(Some(func)) + update_slot!(getattro, func); + } + "__iter__" => { + let func: slots::IterFunc = |zelf, vm| { + let magic = get_class_magic(&zelf, "__iter__"); + vm.invoke(&magic, (zelf,)) + }; + update_slot!(iter, func); + } + "__next__" => { + let func: slots::IterNextFunc = |zelf, vm| { + let magic = get_class_magic(zelf, "__next__"); + vm.invoke(&magic, (zelf.clone(),)) + }; + update_slot!(iternext, func); } _ => {} } @@ -327,10 +347,11 @@ impl PyType { } } let attr_name = attr_name.borrow_value(); + let mut attributes = zelf.attributes.write(); + attributes.insert(attr_name.to_owned(), value); if attr_name.starts_with("__") && attr_name.ends_with("__") { - zelf.update_slot(attr_name); + zelf.update_slot(attr_name, true); } - zelf.attributes.write().insert(attr_name.to_owned(), value); Ok(()) } @@ -342,9 +363,17 @@ impl PyType { } } - zelf.get_attr(attr_name.borrow_value()) - .ok_or_else(|| vm.new_attribute_error(attr_name.borrow_value().to_owned()))?; - zelf.attributes.write().remove(attr_name.borrow_value()); + let mut attributes = zelf.attributes.write(); + if attributes.remove(attr_name.borrow_value()).is_none() { + return Err(vm.new_exception( + vm.ctx.exceptions.attribute_error.clone(), + vec![attr_name.into_object()], + )); + }; + let attr_name = attr_name.borrow_value(); + if attr_name.starts_with("__") && attr_name.ends_with("__") { + zelf.update_slot(attr_name, false); + } Ok(()) } @@ -800,7 +829,7 @@ pub fn new( for attr_name in new_type.attributes.read().keys() { if attr_name.starts_with("__") && attr_name.ends_with("__") { - new_type.update_slot(attr_name); + new_type.update_slot(attr_name, true); } } for base in &new_type.bases { diff --git a/vm/src/builtins/range.rs b/vm/src/builtins/range.rs index 52132cd57..5f257d65b 100644 --- a/vm/src/builtins/range.rs +++ b/vm/src/builtins/range.rs @@ -12,7 +12,7 @@ use crate::pyobject::{ self, BorrowValue, IdProtocol, IntoPyRef, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol, }; -use crate::slots::{Comparable, Hashable, PyComparisonOp}; +use crate::slots::{Comparable, Hashable, Iterable, PyComparisonOp, PyIter}; use crate::vm::VirtualMachine; /// range(stop) -> range object @@ -128,7 +128,7 @@ pub fn init(context: &PyContext) { type PyRangeRef = PyRef; -#[pyimpl(with(Hashable, Comparable))] +#[pyimpl(with(Hashable, Comparable, Iterable))] impl PyRange { fn new(cls: PyTypeRef, stop: PyIntRef, vm: &VirtualMachine) -> PyResult> { PyRange { @@ -168,14 +168,6 @@ impl PyRange { self.step.clone() } - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRangeIterator { - PyRangeIterator { - position: AtomicCell::new(0), - range: zelf, - } - } - #[pymethod(name = "__reversed__")] fn reversed(&self, vm: &VirtualMachine) -> PyRangeIterator { let start = self.start.borrow_value(); @@ -368,6 +360,16 @@ impl Comparable for PyRange { } } +impl Iterable for PyRange { + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + Ok(PyRangeIterator { + position: AtomicCell::new(0), + range: zelf, + } + .into_object(vm)) + } +} + #[pyclass(module = false, name = "range_iterator")] #[derive(Debug)] pub struct PyRangeIterator { @@ -381,24 +383,18 @@ impl PyValue for PyRangeIterator { } } -type PyRangeIteratorRef = PyRef; +#[pyimpl(with(PyIter))] +impl PyRangeIterator {} -#[pyimpl] -impl PyRangeIterator { - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let position = BigInt::from(self.position.fetch_add(1)); - if let Some(int) = self.range.get(&position) { - Ok(int) +impl PyIter for PyRangeIterator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let position = BigInt::from(zelf.position.fetch_add(1)); + if let Some(int) = zelf.range.get(&position) { + Ok(vm.ctx.new_int(int)) } else { Err(vm.new_stop_iteration()) } } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRangeIteratorRef { - zelf - } } pub enum RangeIndex { diff --git a/vm/src/builtins/set.rs b/vm/src/builtins/set.rs index ccea079ff..cdb7953a4 100644 --- a/vm/src/builtins/set.rs +++ b/vm/src/builtins/set.rs @@ -10,8 +10,9 @@ use crate::pyobject::{ self, BorrowValue, PyClassImpl, PyComparisonValue, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol, }; -use crate::slots::{Comparable, Hashable, PyComparisonOp, Unhashable}; +use crate::slots::{Comparable, Hashable, Iterable, PyComparisonOp, PyIter, Unhashable}; use crate::vm::{ReprGuard, VirtualMachine}; +use crossbeam_utils::atomic::AtomicCell; use std::fmt; pub type SetContentType = dictdatatype::Dict<()>; @@ -192,7 +193,7 @@ impl PySetInner { Ok(true) } - fn iter(&self, _vm: &VirtualMachine) -> PySetIterator { + fn iter(&self) -> PySetIterator { let set_size = SetSizeInfo { position: 0, size: Some(self.content.len()), @@ -200,7 +201,7 @@ impl PySetInner { PySetIterator { dict: PyRc::clone(&self.content), - size_info: crossbeam_utils::atomic::AtomicCell::new(set_size), + size_info: AtomicCell::new(set_size), } } @@ -311,7 +312,7 @@ macro_rules! multi_args_set { }}; } -#[pyimpl(with(Hashable, Comparable), flags(BASETYPE))] +#[pyimpl(with(Hashable, Comparable, Iterable), flags(BASETYPE))] impl PySet { #[pyslot] fn tp_new( @@ -426,11 +427,6 @@ impl PySet { self.xor(other, vm) } - #[pymethod(name = "__iter__")] - fn iter(&self, vm: &VirtualMachine) -> PySetIterator { - self.inner.iter(vm) - } - #[pymethod(name = "__repr__")] fn repr(zelf: PyRef, vm: &VirtualMachine) -> PyResult { let s = if zelf.inner.len() == 0 { @@ -539,6 +535,12 @@ impl Comparable for PySet { impl Unhashable for PySet {} +impl Iterable for PySet { + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + Ok(zelf.inner.iter().into_object(vm)) + } +} + macro_rules! multi_args_frozenset { ($vm:expr, $others:expr, $zelf:expr, $op:tt) => {{ let mut res = $zelf.inner.copy(); @@ -549,7 +551,7 @@ macro_rules! multi_args_frozenset { }}; } -#[pyimpl(flags(BASETYPE), with(Hashable, Comparable))] +#[pyimpl(flags(BASETYPE), with(Hashable, Comparable, Iterable))] impl PyFrozenSet { // used by ssl.rs windows #[allow(dead_code)] @@ -677,11 +679,6 @@ impl PyFrozenSet { self.xor(other, vm) } - #[pymethod(name = "__iter__")] - fn iter(&self, vm: &VirtualMachine) -> PySetIterator { - self.inner.iter(vm) - } - #[pymethod(name = "__repr__")] fn repr(zelf: PyRef, vm: &VirtualMachine) -> PyResult { let inner = &zelf.inner; @@ -715,6 +712,12 @@ impl Comparable for PyFrozenSet { } } +impl Iterable for PyFrozenSet { + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + Ok(zelf.inner.iter().into_object(vm)) + } +} + struct SetIterable { iterable: Args, } @@ -745,7 +748,7 @@ struct SetSizeInfo { #[pyclass(module = false, name = "set_iterator")] pub(crate) struct PySetIterator { dict: PyRc, - size_info: crossbeam_utils::atomic::AtomicCell, + size_info: AtomicCell, } impl fmt::Debug for PySetIterator { @@ -755,34 +758,14 @@ impl fmt::Debug for PySetIterator { } } -#[pyimpl] +impl PyValue for PySetIterator { + fn class(vm: &VirtualMachine) -> &PyTypeRef { + &vm.ctx.types.set_iterator_type + } +} + +#[pyimpl(with(PyIter))] impl PySetIterator { - #[pymethod(magic)] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let mut size_info = self.size_info.take(); - - if let Some(set_size) = size_info.size { - if set_size == self.dict.len() { - let index = size_info.position; - let keys = self.dict.keys(); - let item = keys.get(index).ok_or_else(|| vm.new_stop_iteration())?; - size_info.position += 1; - self.size_info.store(size_info); - return Ok(item.clone()); - } else { - size_info.size = None; - self.size_info.store(size_info); - } - } - - Err(vm.new_runtime_error("set changed size during iteration".into())) - } - - #[pymethod(magic)] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - #[pymethod(magic)] fn length_hint(&self) -> usize { let set_len = self.dict.len(); @@ -791,9 +774,25 @@ impl PySetIterator { } } -impl PyValue for PySetIterator { - fn class(vm: &VirtualMachine) -> &PyTypeRef { - &vm.ctx.types.set_iterator_type +impl PyIter for PySetIterator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let mut size_info = zelf.size_info.take(); + + if let Some(set_size) = size_info.size { + if set_size == zelf.dict.len() { + let index = size_info.position; + let keys = zelf.dict.keys(); + let item = keys.get(index).ok_or_else(|| vm.new_stop_iteration())?; + size_info.position += 1; + zelf.size_info.store(size_info); + return Ok(item.clone()); + } else { + size_info.size = None; + zelf.size_info.store(size_info); + } + } + + Err(vm.new_runtime_error("set changed size during iteration".into())) } } diff --git a/vm/src/builtins/tuple.rs b/vm/src/builtins/tuple.rs index e88f378b3..91ce4050d 100644 --- a/vm/src/builtins/tuple.rs +++ b/vm/src/builtins/tuple.rs @@ -10,7 +10,7 @@ use crate::pyobject::{ }; use crate::sequence::{self, SimpleSeq}; use crate::sliceable::PySliceableSequence; -use crate::slots::{Comparable, Hashable, PyComparisonOp}; +use crate::slots::{Comparable, Hashable, Iterable, PyComparisonOp, PyIter}; use crate::vm::{ReprGuard, VirtualMachine}; /// tuple() -> empty tuple @@ -84,7 +84,7 @@ pub(crate) fn get_value(obj: &PyObjectRef) -> &[PyObjectRef] { obj.payload::().unwrap().borrow_value() } -#[pyimpl(flags(BASETYPE), with(Hashable, Comparable))] +#[pyimpl(flags(BASETYPE), with(Hashable, Comparable, Iterable))] impl PyTuple { /// Creating a new tuple with given boxed slice. /// NOTE: for usual case, you probably want to use PyTupleRef::with_elements. @@ -134,14 +134,6 @@ impl PyTuple { Ok(count) } - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyTupleIterator { - PyTupleIterator { - position: AtomicCell::new(0), - tuple: zelf, - } - } - #[pymethod(name = "__len__")] pub(crate) fn len(&self) -> usize { self.elements.len() @@ -258,6 +250,16 @@ impl Comparable for PyTuple { } } +impl Iterable for PyTuple { + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + Ok(PyTupleIterator { + position: AtomicCell::new(0), + tuple: zelf, + } + .into_object(vm)) + } +} + #[pyclass(module = false, name = "tuple_iterator")] #[derive(Debug)] pub struct PyTupleIterator { @@ -271,22 +273,18 @@ impl PyValue for PyTupleIterator { } } -#[pyimpl] -impl PyTupleIterator { - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let pos = self.position.fetch_add(1); - if let Some(obj) = self.tuple.borrow_value().get(pos) { +#[pyimpl(with(PyIter))] +impl PyTupleIterator {} + +impl PyIter for PyTupleIterator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let pos = zelf.position.fetch_add(1); + if let Some(obj) = zelf.tuple.borrow_value().get(pos) { Ok(obj.clone()) } else { Err(vm.new_stop_iteration()) } } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } pub fn init(context: &PyContext) { diff --git a/vm/src/builtins/zip.rs b/vm/src/builtins/zip.rs index 9791a3af7..e236c0c47 100644 --- a/vm/src/builtins/zip.rs +++ b/vm/src/builtins/zip.rs @@ -2,6 +2,7 @@ use super::pytype::PyTypeRef; use crate::function::Args; use crate::iterator; use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; +use crate::slots::PyIter; use crate::vm::VirtualMachine; pub type PyZipRef = PyRef; @@ -18,23 +19,24 @@ impl PyValue for PyZip { } } -#[pyimpl(flags(BASETYPE))] +#[pyimpl(with(PyIter), flags(BASETYPE))] impl PyZip { #[pyslot] fn tp_new(cls: PyTypeRef, iterables: Args, vm: &VirtualMachine) -> PyResult { let iterators = iterables .into_iter() - .map(|iterable| iterator::get_iter(vm, &iterable)) + .map(|iterable| iterator::get_iter(vm, iterable)) .collect::, _>>()?; PyZip { iterators }.into_ref_with_type(vm, cls) } +} - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - if self.iterators.is_empty() { +impl PyIter for PyZip { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + if zelf.iterators.is_empty() { Err(vm.new_stop_iteration()) } else { - let next_objs = self + let next_objs = zelf .iterators .iter() .map(|iterator| iterator::call_next(vm, iterator)) @@ -43,11 +45,6 @@ impl PyZip { Ok(vm.ctx.new_tuple(next_objs)) } } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } pub fn init(context: &PyContext) { diff --git a/vm/src/frame.rs b/vm/src/frame.rs index 479f86076..aca4c3515 100644 --- a/vm/src/frame.rs +++ b/vm/src/frame.rs @@ -583,7 +583,7 @@ impl ExecutingFrame<'_> { } bytecode::Instruction::GetIter => { let iterated_obj = self.pop_value(); - let iter_obj = iterator::get_iter(vm, &iterated_obj)?; + let iter_obj = iterator::get_iter(vm, iterated_obj)?; self.push_value(iter_obj); Ok(None) } diff --git a/vm/src/iterator.rs b/vm/src/iterator.rs index e7040c42f..0eb9265f4 100644 --- a/vm/src/iterator.rs +++ b/vm/src/iterator.rs @@ -14,30 +14,41 @@ use num_traits::Signed; * in the vm when a for loop is entered. Next, it is used when the builtin * function 'iter' is called. */ -pub fn get_iter(vm: &VirtualMachine, iter_target: &PyObjectRef) -> PyResult { - if let Some(method_or_err) = vm.get_method(iter_target.clone(), "__iter__") { - let method = method_or_err?; - let iter = vm.invoke(&method, ())?; - if iter.has_class_attr("__next__") { +pub fn get_iter(vm: &VirtualMachine, iter_target: PyObjectRef) -> PyResult { + let getiter = { + let cls = iter_target.class(); + cls.mro_find_map(|x| x.slots.iter.load()) + }; + if let Some(getiter) = getiter { + let iter = getiter(iter_target, vm)?; + let cls = iter.class(); + let is_iter = cls.iter_mro().any(|x| x.slots.iternext.load().is_some()); + if is_iter { + drop(cls); Ok(iter) } else { Err(vm.new_type_error(format!( "iter() returned non-iterator of type '{}'", - iter.class().name + cls.name ))) } } else { vm.get_method_or_type_error(iter_target.clone(), "__getitem__", || { format!("'{}' object is not iterable", iter_target.class().name) })?; - Ok(PySequenceIterator::new_forward(iter_target.clone()) + Ok(PySequenceIterator::new_forward(iter_target) .into_ref(vm) .into_object()) } } pub fn call_next(vm: &VirtualMachine, iter_obj: &PyObjectRef) -> PyResult { - vm.call_method(iter_obj, "__next__", ()) + let iternext = { + let cls = iter_obj.class(); + cls.mro_find_map(|x| x.slots.iternext.load()) + .ok_or_else(|| vm.new_type_error(format!("'{}' object is not an iterator", cls.name)))? + }; + iternext(iter_obj, vm) } /* diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs index 116eddd0f..716bcd6bd 100644 --- a/vm/src/pyobject.rs +++ b/vm/src/pyobject.rs @@ -840,25 +840,21 @@ impl fmt::Debug for PyObject { /// PyIterable can optionally perform type checking and conversions on iterated /// objects using a generic type parameter that implements `TryFromObject`. pub struct PyIterable { - method: PyObjectRef, - _item: std::marker::PhantomData, + iterable: PyObjectRef, + iterfn: Option, + _item: PhantomData, } impl PyIterable { - pub fn from_method(method: PyObjectRef) -> Self { - PyIterable { - method, - _item: std::marker::PhantomData, - } - } - /// Returns an iterator over this sequence of objects. /// /// This operation may fail if an exception is raised while invoking the /// `__iter__` method of the iterable object. pub fn iter<'a>(&self, vm: &'a VirtualMachine) -> PyResult> { - let method = &self.method; - let iter_obj = vm.invoke(method, ())?; + let iter_obj = match self.iterfn { + Some(f) => f(self.iterable.clone(), vm)?, + None => PySequenceIterator::new_forward(self.iterable.clone()).into_object(vm), + }; let length_hint = iterator::length_hint(vm, iter_obj.clone())?; @@ -866,7 +862,7 @@ impl PyIterable { vm, obj: iter_obj, length_hint, - _item: std::marker::PhantomData, + _item: PhantomData, }) } } @@ -876,23 +872,19 @@ where T: TryFromObject, { fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { - if let Some(method_or_err) = vm.get_method(obj.clone(), "__iter__") { - let method = method_or_err?; - Ok(PyIterable { - method, - _item: std::marker::PhantomData, - }) - } else { - vm.get_method_or_type_error(obj.clone(), "__getitem__", || { - format!("'{}' object is not iterable", obj.class().name) - })?; - Self::try_from_object( - vm, - PySequenceIterator::new_forward(obj) - .into_ref(vm) - .into_object(), - ) + let iterfn; + { + let cls = obj.class(); + iterfn = cls.mro_find_map(|x| x.slots.iter.load()); + if iterfn.is_none() && !cls.has_attr("__getitem__") { + return Err(vm.new_type_error(format!("'{}' object is not iterable", cls.name))); + } } + Ok(PyIterable { + iterable: obj, + iterfn, + _item: PhantomData, + }) } } @@ -900,7 +892,7 @@ pub struct PyIterator<'a, T> { vm: &'a VirtualMachine, obj: PyObjectRef, length_hint: Option, - _item: std::marker::PhantomData, + _item: PhantomData, } impl<'a, T> Iterator for PyIterator<'a, T> @@ -910,16 +902,9 @@ where type Item = PyResult; fn next(&mut self) -> Option { - match self.vm.call_method(&self.obj, "__next__", ()) { - Ok(value) => Some(T::try_from_object(self.vm, value)), - Err(err) => { - if err.isinstance(&self.vm.ctx.exceptions.stop_iteration) { - None - } else { - Some(Err(err)) - } - } - } + iterator::get_next_object(self.vm, &self.obj) + .transpose() + .map(|x| x.and_then(|obj| T::try_from_object(self.vm, obj))) } fn size_hint(&self) -> (usize, Option) { diff --git a/vm/src/slots.rs b/vm/src/slots.rs index 711e9feaf..833f271d9 100644 --- a/vm/src/slots.rs +++ b/vm/src/slots.rs @@ -55,6 +55,8 @@ pub(crate) type CmpFunc = fn( ) -> PyResult>; pub(crate) type GetattroFunc = fn(PyObjectRef, PyStrRef, &VirtualMachine) -> PyResult; pub(crate) type BufferFunc = fn(&PyObjectRef, &VirtualMachine) -> PyResult>; +pub(crate) type IterFunc = fn(PyObjectRef, &VirtualMachine) -> PyResult; +pub(crate) type IterNextFunc = fn(&PyObjectRef, &VirtualMachine) -> PyResult; #[derive(Default)] pub struct PyTypeSlots { @@ -68,6 +70,8 @@ pub struct PyTypeSlots { pub cmp: AtomicCell>, pub getattro: AtomicCell>, pub buffer: Option, + pub iter: AtomicCell>, + pub iternext: AtomicCell>, } impl PyTypeSlots { @@ -418,3 +422,49 @@ pub trait BufferProtocol: PyValue { fn get_buffer(zelf: &PyRef, vm: &VirtualMachine) -> PyResult>; } + +#[pyimpl] +pub trait Iterable: PyValue { + #[pyslot] + fn tp_iter(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult { + if let Ok(zelf) = zelf.downcast() { + Self::iter(zelf, vm) + } else { + Err(vm.new_type_error("unexpected payload for __iter__".to_owned())) + } + } + + #[pymethod(magic)] + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult; +} + +#[pyimpl(with(Iterable))] +pub trait PyIter: PyValue { + #[pyslot] + fn tp_iternext(zelf: &PyObjectRef, vm: &VirtualMachine) -> PyResult { + if let Some(zelf) = zelf.downcast_ref() { + Self::next(zelf, vm) + } else { + Err(vm.new_type_error("unexpected payload for __next__".to_owned())) + } + } + + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult; + + #[pymethod] + fn __next__(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + Self::next(&zelf, vm) + } +} + +impl Iterable for T +where + T: PyIter, +{ + fn tp_iter(zelf: PyObjectRef, _vm: &VirtualMachine) -> PyResult { + Ok(zelf) + } + fn iter(zelf: PyRef, _vm: &VirtualMachine) -> PyResult { + Ok(zelf.into_object()) + } +} diff --git a/vm/src/stdlib/array.rs b/vm/src/stdlib/array.rs index 2b33d2736..5530d7ae1 100644 --- a/vm/src/stdlib/array.rs +++ b/vm/src/stdlib/array.rs @@ -16,7 +16,7 @@ use crate::pyobject::{ PyObjectRef, PyRef, PyResult, PyValue, StaticType, TryFromObject, TypeProtocol, }; use crate::sliceable::{saturate_index, PySliceableSequence, PySliceableSequenceMut}; -use crate::slots::{BufferProtocol, Comparable, PyComparisonOp}; +use crate::slots::{BufferProtocol, Comparable, Iterable, PyComparisonOp, PyIter}; use crate::VirtualMachine; use crossbeam_utils::atomic::AtomicCell; use itertools::Itertools; @@ -486,7 +486,7 @@ impl From for PyArray { } } -#[pyimpl(flags(BASETYPE), with(Comparable, BufferProtocol))] +#[pyimpl(flags(BASETYPE), with(Comparable, BufferProtocol, Iterable))] impl PyArray { fn borrow_value(&self) -> PyRwLockReadGuard<'_, ArrayContentType> { self.array.read() @@ -766,14 +766,6 @@ impl PyArray { self.borrow_value().len() } - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyArrayIter { - PyArrayIter { - position: AtomicCell::new(0), - array: zelf, - } - } - fn array_eq(&self, other: &Self, vm: &VirtualMachine) -> PyResult { // we cannot use zelf.is(other) for shortcut because if we contenting a // float value NaN we always return False even they are the same object. @@ -889,6 +881,16 @@ impl Buffer for ArrayBuffer { } } +impl Iterable for PyArray { + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + Ok(PyArrayIter { + position: AtomicCell::new(0), + array: zelf, + } + .into_object(vm)) + } +} + impl<'a> ResizeGuard<'a> for PyArray { type Resizable = PyRwLockWriteGuard<'a, ArrayContentType>; @@ -916,22 +918,18 @@ impl PyValue for PyArrayIter { } } -#[pyimpl] -impl PyArrayIter { - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let pos = self.position.fetch_add(1); - if let Some(item) = self.array.borrow_value().getitem_by_idx(pos, vm) { +#[pyimpl(with(PyIter))] +impl PyArrayIter {} + +impl PyIter for PyArrayIter { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let pos = zelf.position.fetch_add(1); + if let Some(item) = zelf.array.borrow_value().getitem_by_idx(pos, vm) { Ok(item) } else { Err(vm.new_stop_iteration()) } } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { diff --git a/vm/src/stdlib/collections.rs b/vm/src/stdlib/collections.rs index 59f088674..ddddc8e58 100644 --- a/vm/src/stdlib/collections.rs +++ b/vm/src/stdlib/collections.rs @@ -8,7 +8,7 @@ mod _collections { use crate::pyobject::{ PyComparisonValue, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, StaticType, }; - use crate::slots::{Comparable, PyComparisonOp}; + use crate::slots::{Comparable, Iterable, PyComparisonOp, PyIter}; use crate::vm::ReprGuard; use crate::VirtualMachine; use crate::{sequence, sliceable}; @@ -66,7 +66,7 @@ mod _collections { } } - #[pyimpl(flags(BASETYPE), with(Comparable))] + #[pyimpl(flags(BASETYPE), with(Comparable, Iterable))] impl PyDeque { #[pyslot] fn tp_new( @@ -323,14 +323,6 @@ mod _collections { fn len(&self) -> usize { self.borrow_deque().len() } - - #[pymethod(magic)] - fn iter(zelf: PyRef) -> PyDequeIterator { - PyDequeIterator { - position: AtomicCell::new(0), - deque: zelf, - } - } } impl Comparable for PyDeque { @@ -350,6 +342,16 @@ mod _collections { } } + impl Iterable for PyDeque { + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + Ok(PyDequeIterator { + position: AtomicCell::new(0), + deque: zelf, + } + .into_object(vm)) + } + } + #[pyattr] #[pyclass(name = "_deque_iterator")] #[derive(Debug)] @@ -364,12 +366,13 @@ mod _collections { } } - #[pyimpl] - impl PyDequeIterator { - #[pymethod(magic)] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let pos = self.position.fetch_add(1); - let deque = self.deque.borrow_deque(); + #[pyimpl(with(PyIter))] + impl PyDequeIterator {} + + impl PyIter for PyDequeIterator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let pos = zelf.position.fetch_add(1); + let deque = zelf.deque.borrow_deque(); if pos < deque.len() { let ret = deque[pos].clone(); Ok(ret) @@ -377,10 +380,5 @@ mod _collections { Err(vm.new_stop_iteration()) } } - - #[pymethod(magic)] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } } diff --git a/vm/src/stdlib/csv.rs b/vm/src/stdlib/csv.rs index 10c06717c..9ffcd2088 100644 --- a/vm/src/stdlib/csv.rs +++ b/vm/src/stdlib/csv.rs @@ -7,8 +7,8 @@ use crate::builtins::pytype::PyTypeRef; use crate::common::lock::PyRwLock; use crate::function::FuncArgs; use crate::pyobject::{ - BorrowValue, IntoPyObject, PyClassImpl, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, - StaticType, TryFromObject, TypeProtocol, + BorrowValue, IntoPyObject, PyClassImpl, PyIterable, PyObjectRef, PyResult, PyValue, StaticType, + TryFromObject, TypeProtocol, }; use crate::types::create_simple_type; use crate::VirtualMachine; @@ -149,15 +149,24 @@ impl Reader { #[pyimpl] impl Reader { - #[pymethod(name = "__iter__")] - fn iter(this: PyRef, vm: &VirtualMachine) -> PyResult { + #[pyslot] + #[pymethod(magic)] + fn iter(this: PyObjectRef, vm: &VirtualMachine) -> PyResult { + let this = match this.downcast::() { + Ok(reader) => reader, + Err(_) => return Err(vm.new_type_error("unexpected payload for __iter__".to_owned())), + }; this.state.write().cast_to_reader(vm)?; Ok(this.into_pyobject(vm)) } - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let mut state = self.state.write(); + #[pyslot] + fn tp_iternext(zelf: &PyObjectRef, vm: &VirtualMachine) -> PyResult { + let zelf = match zelf.downcast_ref::() { + Some(reader) => reader, + None => return Err(vm.new_type_error("unexpected payload for __next__".to_owned())), + }; + let mut state = zelf.state.write(); state.cast_to_reader(vm)?; if let ReadState::CsvIter(ref mut reader) = &mut *state { diff --git a/vm/src/stdlib/functools.rs b/vm/src/stdlib/functools.rs index 5f9cbba7b..2afe6da85 100644 --- a/vm/src/stdlib/functools.rs +++ b/vm/src/stdlib/functools.rs @@ -14,7 +14,7 @@ mod _functools { start_value: OptionalArg, vm: &VirtualMachine, ) -> PyResult { - let iterator = iterator::get_iter(vm, &sequence)?; + let iterator = iterator::get_iter(vm, sequence)?; let start_value = if let OptionalArg::Present(val) = start_value { val diff --git a/vm/src/stdlib/io.rs b/vm/src/stdlib/io.rs index 154e9035d..0595ed027 100644 --- a/vm/src/stdlib/io.rs +++ b/vm/src/stdlib/io.rs @@ -453,13 +453,14 @@ mod _io { check_seekable(&instance, vm) } - #[pymethod(magic)] - fn iter(instance: PyObjectRef, vm: &VirtualMachine) -> PyResult { + #[pyslot] + #[pymethod(name = "__iter__")] + fn tp_iter(instance: PyObjectRef, vm: &VirtualMachine) -> PyResult { check_closed(&instance, vm)?; Ok(instance) } - #[pymethod(magic)] - fn next(instance: PyObjectRef, vm: &VirtualMachine) -> PyResult { + #[pyslot] + fn tp_iternext(instance: &PyObjectRef, vm: &VirtualMachine) -> PyResult { let line = call_method(vm, &instance, "readline", ())?; if !pybool::boolval(vm, line.clone())? { Err(vm.new_stop_iteration()) @@ -467,6 +468,10 @@ mod _io { Ok(line) } } + #[pymethod(magic)] + fn next(instance: PyObjectRef, vm: &VirtualMachine) -> PyResult { + Self::tp_iternext(&instance, vm) + } } pub(super) fn iobase_close(file: &PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index a2b886066..817fb124d 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -16,9 +16,10 @@ mod decl { use crate::function::{Args, FuncArgs, OptionalArg, OptionalOption}; use crate::iterator::{call_next, get_all, get_iter, get_next_object}; use crate::pyobject::{ - BorrowValue, IdProtocol, IntoPyRef, PyCallable, PyObjectRc, PyObjectRef, PyObjectWeak, + BorrowValue, IdProtocol, IntoPyObject, PyCallable, PyObjectRc, PyObjectRef, PyObjectWeak, PyRef, PyResult, PyValue, StaticType, TypeProtocol, }; + use crate::slots::PyIter; use crate::vm::VirtualMachine; #[pyattr] @@ -36,7 +37,7 @@ mod decl { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsChain { #[pyslot] fn tp_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult> { @@ -48,19 +49,36 @@ mod decl { .into_ref_with_type(vm, cls) } - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { + #[pyclassmethod(name = "from_iterable")] + fn from_iterable( + cls: PyTypeRef, + iterable: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult> { + let it = get_iter(vm, iterable)?; + let iterables = get_all(vm, &it)?; + + PyItertoolsChain { + iterables, + cur_idx: AtomicCell::new(0), + cached_iter: PyRwLock::new(None), + } + .into_ref_with_type(vm, cls) + } + } + impl PyIter for PyItertoolsChain { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { loop { - let pos = self.cur_idx.load(); - if pos >= self.iterables.len() { + let pos = zelf.cur_idx.load(); + if pos >= zelf.iterables.len() { break; } - let cur_iter = if self.cached_iter.read().is_none() { + let cur_iter = if zelf.cached_iter.read().is_none() { // We need to call "get_iter" outside of the lock. - let iter = get_iter(vm, &self.iterables[pos])?; - *self.cached_iter.write() = Some(iter.clone()); + let iter = get_iter(vm, zelf.iterables[pos].clone())?; + *zelf.cached_iter.write() = Some(iter.clone()); iter - } else if let Some(cached_iter) = (*(self.cached_iter.read())).clone() { + } else if let Some(cached_iter) = (*zelf.cached_iter.read()).clone() { cached_iter } else { // Someone changed cached iter to None since we checked. @@ -72,8 +90,8 @@ mod decl { Ok(ok) => return Ok(ok), Err(err) => { if err.isinstance(&vm.ctx.exceptions.stop_iteration) { - self.cur_idx.fetch_add(1); - *self.cached_iter.write() = None; + zelf.cur_idx.fetch_add(1); + *zelf.cached_iter.write() = None; } else { return Err(err); } @@ -83,28 +101,6 @@ mod decl { Err(vm.new_stop_iteration()) } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - - #[pyclassmethod(name = "from_iterable")] - fn from_iterable( - cls: PyTypeRef, - iterable: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult> { - let it = get_iter(vm, &iterable)?; - let iterables = get_all(vm, &it)?; - - PyItertoolsChain { - iterables, - cur_idx: AtomicCell::new(0), - cached_iter: PyRwLock::new(None), - } - .into_ref_with_type(vm, cls) - } } #[pyattr] @@ -121,7 +117,7 @@ mod decl { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsCompress { #[pyslot] fn tp_new( @@ -130,8 +126,8 @@ mod decl { selector: PyObjectRef, vm: &VirtualMachine, ) -> PyResult> { - let data_iter = get_iter(vm, &data)?; - let selector_iter = get_iter(vm, &selector)?; + let data_iter = get_iter(vm, data)?; + let selector_iter = get_iter(vm, selector)?; PyItertoolsCompress { data: data_iter, @@ -139,24 +135,19 @@ mod decl { } .into_ref_with_type(vm, cls) } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { + } + impl PyIter for PyItertoolsCompress { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { loop { - let sel_obj = call_next(vm, &self.selector)?; + let sel_obj = call_next(vm, &zelf.selector)?; let verdict = pybool::boolval(vm, sel_obj.clone())?; - let data_obj = call_next(vm, &self.data)?; + let data_obj = call_next(vm, &zelf.data)?; if verdict { return Ok(data_obj); } } } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } #[pyattr] @@ -173,7 +164,7 @@ mod decl { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsCount { #[pyslot] fn tp_new( @@ -197,18 +188,13 @@ mod decl { } .into_ref_with_type(vm, cls) } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let mut cur = self.cur.write(); + } + impl PyIter for PyItertoolsCount { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let mut cur = zelf.cur.write(); let result = cur.clone(); - *cur += &self.step; - Ok(result.into_pyref(vm)) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf + *cur += &zelf.step; + Ok(result.into_pyobject(vm)) } } @@ -227,7 +213,7 @@ mod decl { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsCycle { #[pyslot] fn tp_new( @@ -235,7 +221,7 @@ mod decl { iterable: PyObjectRef, vm: &VirtualMachine, ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; + let iter = get_iter(vm, iterable)?; PyItertoolsCycle { iter, @@ -244,22 +230,22 @@ mod decl { } .into_ref_with_type(vm, cls) } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let item = if let Some(item) = get_next_object(vm, &self.iter)? { - self.saved.write().push(item.clone()); + } + impl PyIter for PyItertoolsCycle { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let item = if let Some(item) = get_next_object(vm, &zelf.iter)? { + zelf.saved.write().push(item.clone()); item } else { - let saved = self.saved.read(); + let saved = zelf.saved.read(); if saved.len() == 0 { return Err(vm.new_stop_iteration()); } - let last_index = self.index.fetch_add(1); + let last_index = zelf.index.fetch_add(1); if last_index >= saved.len() - 1 { - self.index.store(0); + zelf.index.store(0); } saved[last_index].clone() @@ -267,11 +253,6 @@ mod decl { Ok(item) } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } #[pyattr] @@ -288,7 +269,7 @@ mod decl { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsRepeat { #[pyslot] fn tp_new( @@ -305,9 +286,17 @@ mod decl { PyItertoolsRepeat { object, times }.into_ref_with_type(vm, cls) } - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - if let Some(ref times) = self.times { + #[pymethod(name = "__length_hint__")] + fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef { + match self.times { + Some(ref times) => vm.ctx.new_int(times.read().clone()), + None => vm.ctx.new_int(0), + } + } + } + impl PyIter for PyItertoolsRepeat { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + if let Some(ref times) = zelf.times { let mut times = times.write(); if !times.is_positive() { return Err(vm.new_stop_iteration()); @@ -315,20 +304,7 @@ mod decl { *times -= 1; } - Ok(self.object.clone()) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - - #[pymethod(name = "__length_hint__")] - fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef { - match self.times { - Some(ref times) => vm.ctx.new_int(times.read().clone()), - None => vm.ctx.new_int(0), - } + Ok(zelf.object.clone()) } } @@ -346,7 +322,7 @@ mod decl { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsStarmap { #[pyslot] fn tp_new( @@ -355,23 +331,18 @@ mod decl { iterable: PyObjectRef, vm: &VirtualMachine, ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; + let iter = get_iter(vm, iterable)?; PyItertoolsStarmap { function, iter }.into_ref_with_type(vm, cls) } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let obj = call_next(vm, &self.iter)?; - let function = &self.function; + } + impl PyIter for PyItertoolsStarmap { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let obj = call_next(vm, &zelf.iter)?; + let function = &zelf.function; vm.invoke(function, vm.extract_elements(&obj)?) } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } #[pyattr] @@ -389,7 +360,7 @@ mod decl { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsTakewhile { #[pyslot] fn tp_new( @@ -398,7 +369,7 @@ mod decl { iterable: PyObjectRef, vm: &VirtualMachine, ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; + let iter = get_iter(vm, iterable)?; PyItertoolsTakewhile { predicate, @@ -407,31 +378,26 @@ mod decl { } .into_ref_with_type(vm, cls) } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - if self.stop_flag.load() { + } + impl PyIter for PyItertoolsTakewhile { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + if zelf.stop_flag.load() { return Err(vm.new_stop_iteration()); } // might be StopIteration or anything else, which is propagated upwards - let obj = call_next(vm, &self.iterable)?; - let predicate = &self.predicate; + let obj = call_next(vm, &zelf.iterable)?; + let predicate = &zelf.predicate; let verdict = vm.invoke(predicate, (obj.clone(),))?; let verdict = pybool::boolval(vm, verdict)?; if verdict { Ok(obj) } else { - self.stop_flag.store(true); + zelf.stop_flag.store(true); Err(vm.new_stop_iteration()) } } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } #[pyattr] @@ -449,7 +415,7 @@ mod decl { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsDropwhile { #[pyslot] fn tp_new( @@ -458,7 +424,7 @@ mod decl { iterable: PyObjectRef, vm: &VirtualMachine, ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; + let iter = get_iter(vm, iterable)?; PyItertoolsDropwhile { predicate, @@ -467,30 +433,25 @@ mod decl { } .into_ref_with_type(vm, cls) } + } + impl PyIter for PyItertoolsDropwhile { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let predicate = &zelf.predicate; + let iterable = &zelf.iterable; - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let predicate = &self.predicate; - let iterable = &self.iterable; - - if !self.start_flag.load() { + if !zelf.start_flag.load() { loop { let obj = call_next(vm, iterable)?; let pred = predicate.clone(); let pred_value = vm.invoke(&pred.into_object(), (obj.clone(),))?; if !pybool::boolval(vm, pred_value)? { - self.start_flag.store(true); + zelf.start_flag.store(true); return Ok(obj); } } } call_next(vm, iterable) } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } struct GroupByState { @@ -550,11 +511,11 @@ mod decl { key: OptionalOption, } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsGroupBy { #[pyslot] fn tp_new(cls: PyTypeRef, args: GroupByArgs, vm: &VirtualMachine) -> PyResult> { - let iter = get_iter(vm, &args.iterable)?; + let iter = get_iter(vm, args.iterable)?; PyItertoolsGroupBy { iterable: iter, @@ -569,11 +530,18 @@ mod decl { .into_ref_with_type(vm, cls) } - #[pymethod(name = "__next__")] - fn next( - zelf: PyRef, - vm: &VirtualMachine, - ) -> PyResult<(PyObjectRef, PyItertoolsGrouperRef)> { + pub(super) fn advance(&self, vm: &VirtualMachine) -> PyResult<(PyObjectRef, PyObjectRef)> { + let new_value = call_next(vm, &self.iterable)?; + let new_key = if let Some(ref kf) = self.key_func { + vm.invoke(kf, vec![new_value.clone()])? + } else { + new_value.clone() + }; + Ok((new_value, new_key)) + } + } + impl PyIter for PyItertoolsGroupBy { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let mut state = zelf.state.lock(); state.grouper = None; @@ -606,22 +574,7 @@ mod decl { .into_ref(vm); state.grouper = Some(PyObjectRc::downgrade(&grouper.clone().into_typed_pyobj())); - Ok((state.current_key.as_ref().unwrap().clone(), grouper)) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - - pub(super) fn advance(&self, vm: &VirtualMachine) -> PyResult<(PyObjectRef, PyObjectRef)> { - let new_value = call_next(vm, &self.iterable)?; - let new_key = if let Some(ref kf) = self.key_func { - vm.invoke(kf, vec![new_value.clone()])? - } else { - new_value.clone() - }; - Ok((new_value, new_key)) + Ok((state.current_key.as_ref().unwrap().clone(), grouper).into_pyobject(vm)) } } @@ -640,10 +593,10 @@ mod decl { } } - #[pyimpl] - impl PyItertoolsGrouper { - #[pymethod(name = "__next__")] - fn next(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + #[pyimpl(with(PyIter))] + impl PyItertoolsGrouper {} + impl PyIter for PyItertoolsGrouper { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let old_key = { let mut state = zelf.groupby.state.lock(); @@ -670,11 +623,6 @@ mod decl { Err(vm.new_stop_iteration()) } } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } #[pyattr] @@ -703,7 +651,7 @@ mod decl { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsIslice { #[pyslot] fn tp_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult> { @@ -763,7 +711,7 @@ mod decl { None }; - let iter = get_iter(vm, &iter)?; + let iter = get_iter(vm, iter)?; PyItertoolsIslice { iterable: iter, @@ -774,34 +722,29 @@ mod decl { } .into_ref_with_type(vm, cls) } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - while self.cur.load() < self.next.load() { - call_next(vm, &self.iterable)?; - self.cur.fetch_add(1); + } + impl PyIter for PyItertoolsIslice { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + while zelf.cur.load() < zelf.next.load() { + call_next(vm, &zelf.iterable)?; + zelf.cur.fetch_add(1); } - if let Some(stop) = self.stop { - if self.cur.load() >= stop { + if let Some(stop) = zelf.stop { + if zelf.cur.load() >= stop { return Err(vm.new_stop_iteration()); } } - let obj = call_next(vm, &self.iterable)?; - self.cur.fetch_add(1); + let obj = call_next(vm, &zelf.iterable)?; + zelf.cur.fetch_add(1); // TODO is this overflow check required? attempts to copy CPython. - let (next, ovf) = self.next.load().overflowing_add(self.step); - self.next.store(if ovf { self.stop.unwrap() } else { next }); + let (next, ovf) = zelf.next.load().overflowing_add(zelf.step); + zelf.next.store(if ovf { zelf.stop.unwrap() } else { next }); Ok(obj) } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } #[pyattr] @@ -818,7 +761,7 @@ mod decl { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsFilterFalse { #[pyslot] fn tp_new( @@ -827,7 +770,7 @@ mod decl { iterable: PyObjectRef, vm: &VirtualMachine, ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; + let iter = get_iter(vm, iterable)?; PyItertoolsFilterFalse { predicate, @@ -835,11 +778,11 @@ mod decl { } .into_ref_with_type(vm, cls) } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let predicate = &self.predicate; - let iterable = &self.iterable; + } + impl PyIter for PyItertoolsFilterFalse { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let predicate = &zelf.predicate; + let iterable = &zelf.iterable; loop { let obj = call_next(vm, iterable)?; @@ -854,11 +797,6 @@ mod decl { } } } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } #[pyattr] @@ -876,7 +814,7 @@ mod decl { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsAccumulate { #[pyslot] fn tp_new( @@ -885,7 +823,7 @@ mod decl { binop: OptionalArg, vm: &VirtualMachine, ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; + let iter = get_iter(vm, iterable)?; PyItertoolsAccumulate { iterable: iter, @@ -894,33 +832,28 @@ mod decl { } .into_ref_with_type(vm, cls) } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let iterable = &self.iterable; + } + impl PyIter for PyItertoolsAccumulate { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let iterable = &zelf.iterable; let obj = call_next(vm, iterable)?; - let acc_value = self.acc_value.read().clone(); + let acc_value = zelf.acc_value.read().clone(); let next_acc_value = match acc_value { None => obj, Some(value) => { - if vm.is_none(&self.binop) { + if vm.is_none(&zelf.binop) { vm._add(&value, &obj)? } else { - vm.invoke(&self.binop, vec![value, obj])? + vm.invoke(&zelf.binop, vec![value, obj])? } } }; - *self.acc_value.write() = Some(next_acc_value.clone()); + *zelf.acc_value.write() = Some(next_acc_value.clone()); Ok(next_acc_value) } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } #[derive(Debug)] @@ -932,7 +865,7 @@ mod decl { impl PyItertoolsTeeData { fn new(iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult> { Ok(PyRc::new(PyItertoolsTeeData { - iterable: get_iter(vm, &iterable)?, + iterable: get_iter(vm, iterable)?, values: PyRwLock::new(vec![]), })) } @@ -960,11 +893,11 @@ mod decl { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsTee { fn from_iter(iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult { let class = PyItertoolsTee::class(vm); - let it = get_iter(vm, &iterable)?; + let it = get_iter(vm, iterable)?; if it.class().is(PyItertoolsTee::class(vm)) { return vm.call_method(&it, "__copy__", ()); } @@ -1011,18 +944,13 @@ mod decl { .into_ref_with_type(vm, Self::class(vm).clone())? .into_object()) } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let value = self.tee_data.get_item(vm, self.index.load())?; - self.index.fetch_add(1); + } + impl PyIter for PyItertoolsTee { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let value = zelf.tee_data.get_item(vm, zelf.index.load())?; + zelf.index.fetch_add(1); Ok(value) } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } #[pyattr] @@ -1047,7 +975,7 @@ mod decl { repeat: OptionalArg, } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsProduct { #[pyslot] fn tp_new( @@ -1063,7 +991,7 @@ mod decl { let mut pools = Vec::new(); for arg in iterables.into_iter() { - let it = get_iter(vm, &arg)?; + let it = get_iter(vm, arg)?; let pool = get_all(vm, &it)?; pools.push(pool); @@ -1084,35 +1012,6 @@ mod decl { .into_ref_with_type(vm, cls) } - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - // stop signal - if self.stop.load() { - return Err(vm.new_stop_iteration()); - } - - let pools = &self.pools; - - for p in pools { - if p.is_empty() { - return Err(vm.new_stop_iteration()); - } - } - - let idxs = self.idxs.write(); - let res = vm.ctx.new_tuple( - pools - .iter() - .zip(idxs.iter()) - .map(|(pool, idx)| pool[*idx].clone()) - .collect(), - ); - - self.update_idxs(idxs); - - Ok(res) - } - fn update_idxs(&self, mut idxs: PyRwLockWriteGuard<'_, Vec>) { if idxs.len() == 0 { self.stop.store(true); @@ -1135,10 +1034,34 @@ mod decl { self.cur.store(idxs.len() - 1); } } + } + impl PyIter for PyItertoolsProduct { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + // stop signal + if zelf.stop.load() { + return Err(vm.new_stop_iteration()); + } - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf + let pools = &zelf.pools; + + for p in pools { + if p.is_empty() { + return Err(vm.new_stop_iteration()); + } + } + + let idxs = zelf.idxs.write(); + let res = vm.ctx.new_tuple( + pools + .iter() + .zip(idxs.iter()) + .map(|(pool, idx)| pool[*idx].clone()) + .collect(), + ); + + zelf.update_idxs(idxs); + + Ok(res) } } @@ -1158,7 +1081,7 @@ mod decl { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsCombinations { #[pyslot] fn tp_new( @@ -1167,7 +1090,7 @@ mod decl { r: PyIntRef, vm: &VirtualMachine, ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; + let iter = get_iter(vm, iterable)?; let pool = get_all(vm, &iter)?; let r = r.borrow_value(); @@ -1186,36 +1109,31 @@ mod decl { } .into_ref_with_type(vm, cls) } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { + } + impl PyIter for PyItertoolsCombinations { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { // stop signal - if self.exhausted.load() { + if zelf.exhausted.load() { return Err(vm.new_stop_iteration()); } - let n = self.pool.len(); - let r = self.r.load(); + let n = zelf.pool.len(); + let r = zelf.r.load(); if r == 0 { - self.exhausted.store(true); + zelf.exhausted.store(true); return Ok(vm.ctx.new_tuple(vec![])); } let res = vm.ctx.new_tuple( - self.indices + zelf.indices .read() .iter() - .map(|&i| self.pool[i].clone()) + .map(|&i| zelf.pool[i].clone()) .collect(), ); - let mut indices = self.indices.write(); + let mut indices = zelf.indices.write(); // Scan indices right-to-left until finding one that is not at its maximum (i + n - r). let mut idx = r as isize - 1; @@ -1226,7 +1144,7 @@ mod decl { // If no suitable index is found, then the indices are all at // their maximum value and we're done. if idx < 0 { - self.exhausted.store(true); + zelf.exhausted.store(true); } else { // Increment the current index which we know is not at its // maximum. Then move back to the right setting each index @@ -1258,7 +1176,7 @@ mod decl { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsCombinationsWithReplacement { #[pyslot] fn tp_new( @@ -1267,7 +1185,7 @@ mod decl { r: PyIntRef, vm: &VirtualMachine, ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; + let iter = get_iter(vm, iterable)?; let pool = get_all(vm, &iter)?; let r = r.borrow_value(); @@ -1286,32 +1204,27 @@ mod decl { } .into_ref_with_type(vm, cls) } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { + } + impl PyIter for PyItertoolsCombinationsWithReplacement { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { // stop signal - if self.exhausted.load() { + if zelf.exhausted.load() { return Err(vm.new_stop_iteration()); } - let n = self.pool.len(); - let r = self.r.load(); + let n = zelf.pool.len(); + let r = zelf.r.load(); if r == 0 { - self.exhausted.store(true); + zelf.exhausted.store(true); return Ok(vm.ctx.new_tuple(vec![])); } - let mut indices = self.indices.write(); + let mut indices = zelf.indices.write(); let res = vm .ctx - .new_tuple(indices.iter().map(|&i| self.pool[i].clone()).collect()); + .new_tuple(indices.iter().map(|&i| zelf.pool[i].clone()).collect()); // Scan indices right-to-left until finding one that is not at its maximum (i + n - r). let mut idx = r as isize - 1; @@ -1322,7 +1235,7 @@ mod decl { // If no suitable index is found, then the indices are all at // their maximum value and we're done. if idx < 0 { - self.exhausted.store(true); + zelf.exhausted.store(true); } else { let index = indices[idx as usize] + 1; @@ -1355,7 +1268,7 @@ mod decl { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsPermutations { #[pyslot] fn tp_new( @@ -1364,8 +1277,7 @@ mod decl { r: OptionalOption, vm: &VirtualMachine, ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; - let pool = get_all(vm, &iter)?; + let pool = vm.extract_elements(&iterable)?; let n = pool.len(); // If r is not provided, r == n. If provided, r must be a positive integer, or None. @@ -1395,32 +1307,27 @@ mod decl { } .into_ref_with_type(vm, cls) } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { + } + impl PyIter for PyItertoolsPermutations { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { // stop signal - if self.exhausted.load() { + if zelf.exhausted.load() { return Err(vm.new_stop_iteration()); } - let n = self.pool.len(); - let r = self.r.load(); + let n = zelf.pool.len(); + let r = zelf.r.load(); if n == 0 { - self.exhausted.store(true); + zelf.exhausted.store(true); return Ok(vm.ctx.new_tuple(vec![])); } - let mut result = self.result.write(); + let mut result = zelf.result.write(); if let Some(ref mut result) = *result { - let mut indices = self.indices.write(); - let mut cycles = self.cycles.write(); + let mut indices = zelf.indices.write(); + let mut cycles = zelf.cycles.write(); let mut sentinel = false; // Decrement rightmost cycle, moving leftward upon zero rollover @@ -1449,7 +1356,7 @@ mod decl { } } if !sentinel { - self.exhausted.store(true); + zelf.exhausted.store(true); return Err(vm.new_stop_iteration()); } } else { @@ -1462,7 +1369,7 @@ mod decl { .as_ref() .unwrap() .iter() - .map(|&i| self.pool[i].clone()) + .map(|&i| zelf.pool[i].clone()) .collect(), )) } @@ -1488,7 +1395,7 @@ mod decl { fillvalue: OptionalArg, } - #[pyimpl] + #[pyimpl(with(PyIter))] impl PyItertoolsZipLongest { #[pyslot] fn tp_new( @@ -1500,7 +1407,7 @@ mod decl { let fillvalue = args.fillvalue.unwrap_or_none(vm); let iterators = iterables .into_iter() - .map(|iterable| get_iter(vm, &iterable)) + .map(|iterable| get_iter(vm, iterable)) .collect::, _>>()?; PyItertoolsZipLongest { @@ -1509,17 +1416,17 @@ mod decl { } .into_ref_with_type(vm, cls) } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - if self.iterators.is_empty() { + } + impl PyIter for PyItertoolsZipLongest { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + if zelf.iterators.is_empty() { Err(vm.new_stop_iteration()) } else { let mut result: Vec = Vec::new(); - let mut numactive = self.iterators.len(); + let mut numactive = zelf.iterators.len(); - for idx in 0..self.iterators.len() { - let next_obj = match call_next(vm, &self.iterators[idx]) { + for idx in 0..zelf.iterators.len() { + let next_obj = match call_next(vm, &zelf.iterators[idx]) { Ok(obj) => obj, Err(err) => { if !err.isinstance(&vm.ctx.exceptions.stop_iteration) { @@ -1529,7 +1436,7 @@ mod decl { if numactive == 0 { return Err(vm.new_stop_iteration()); } - self.fillvalue.clone() + zelf.fillvalue.clone() } }; result.push(next_obj); @@ -1537,10 +1444,5 @@ mod decl { Ok(vm.ctx.new_tuple(result)) } } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } } } diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs index 11481bc65..03abc1209 100644 --- a/vm/src/stdlib/os.rs +++ b/vm/src/stdlib/os.rs @@ -25,6 +25,7 @@ use crate::pyobject::{ BorrowValue, Either, IntoPyObject, ItemProtocol, PyObjectRef, PyRef, PyResult, PyStructSequence, PyValue, StaticType, TryFromObject, TypeProtocol, }; +use crate::slots::PyIter; use crate::vm::VirtualMachine; // this is basically what CPython has for Py_off_t; windows uses long long @@ -656,41 +657,13 @@ mod _os { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl ScandirIterator { - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - if self.exhausted.load() { - return Err(vm.new_stop_iteration()); - } - - match self.entries.write().next() { - Some(entry) => match entry { - Ok(entry) => Ok(DirEntry { - entry, - mode: self.mode, - } - .into_ref(vm) - .into_object()), - Err(err) => Err(err.into_pyexception(vm)), - }, - None => { - self.exhausted.store(true); - Err(vm.new_stop_iteration()) - } - } - } - #[pymethod] fn close(&self) { self.exhausted.store(true); } - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - #[pymethod(name = "__enter__")] fn enter(zelf: PyRef) -> PyRef { zelf @@ -701,6 +674,29 @@ mod _os { zelf.close() } } + impl PyIter for ScandirIterator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + if zelf.exhausted.load() { + return Err(vm.new_stop_iteration()); + } + + match zelf.entries.write().next() { + Some(entry) => match entry { + Ok(entry) => Ok(DirEntry { + entry, + mode: zelf.mode, + } + .into_ref(vm) + .into_object()), + Err(err) => Err(err.into_pyexception(vm)), + }, + None => { + zelf.exhausted.store(true); + Err(vm.new_stop_iteration()) + } + } + } + } #[pyfunction] fn scandir(path: OptionalArg, vm: &VirtualMachine) -> PyResult { diff --git a/vm/src/stdlib/pystruct.rs b/vm/src/stdlib/pystruct.rs index 302924742..343dd469c 100644 --- a/vm/src/stdlib/pystruct.rs +++ b/vm/src/stdlib/pystruct.rs @@ -35,6 +35,7 @@ pub(crate) mod _struct { BorrowValue, Either, IntoPyObject, PyObjectRef, PyRef, PyResult, PyValue, StaticType, TryFromObject, }; + use crate::slots::PyIter; use crate::VirtualMachine; #[derive(Debug, Copy, Clone, PartialEq)] @@ -832,30 +833,26 @@ pub(crate) mod _struct { } } - #[pyimpl] + #[pyimpl(with(PyIter))] impl UnpackIterator { - #[pymethod(magic)] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let size = self.format_spec.size(); - let offset = self.offset.fetch_add(size); - if offset + size > self.buffer.len() { - Err(vm.new_stop_iteration()) - } else { - self.buffer - .with_ref(|buf| self.format_spec.unpack(&buf[offset..offset + size], vm)) - } - } - - #[pymethod(magic)] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - #[pymethod(magic)] fn length_hint(&self) -> usize { self.buffer.len().saturating_sub(self.offset.load()) / self.format_spec.size() } } + impl PyIter for UnpackIterator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let size = zelf.format_spec.size(); + let offset = zelf.offset.fetch_add(size); + zelf.buffer.with_ref(|buf| { + if let Some(buf) = buf.get(offset..offset + size) { + zelf.format_spec.unpack(buf, vm).map(|x| x.into_object()) + } else { + Err(vm.new_stop_iteration()) + } + }) + } + } #[pyfunction] fn iter_unpack( diff --git a/vm/src/vm.rs b/vm/src/vm.rs index 15228359c..fb1656284 100644 --- a/vm/src/vm.rs +++ b/vm/src/vm.rs @@ -1011,7 +1011,7 @@ impl VirtualMachine { .map(|obj| T::try_from_object(self, obj.clone())) .collect() } else { - let iter = iterator::get_iter(self, value)?; + let iter = iterator::get_iter(self, value.clone())?; iterator::get_all(self, &iter) } } @@ -1049,7 +1049,7 @@ impl VirtualMachine { ref t @ PyTuple => Ok(t.borrow_value().iter().cloned().map(f).collect()), // TODO: put internal iterable type obj => { - let iter = iterator::get_iter(self, obj)?; + let iter = iterator::get_iter(self, obj.clone())?; Ok(iterator::try_map(self, &iter, f)) } }) @@ -1574,7 +1574,7 @@ impl VirtualMachine { // https://docs.python.org/3/reference/expressions.html#membership-test-operations fn _membership_iter_search(&self, haystack: PyObjectRef, needle: PyObjectRef) -> PyResult { - let iter = iterator::get_iter(self, &haystack)?; + let iter = iterator::get_iter(self, haystack)?; loop { if let Some(element) = iterator::get_next_object(self, &iter)? { if self.bool_eq(&needle, &element)? {