Skip to content

Commit 98a62c5

Browse files
committed
Universal mechanism for Iterable/PyIter
1 parent f0dd29d commit 98a62c5

25 files changed

+117
-43
lines changed

vm/src/builtins/asyncgenerator.rs

+3-1
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use crate::coroutine::{Coro, Variant};
55
use crate::exceptions::PyBaseExceptionRef;
66
use crate::frame::FrameRef;
77
use crate::function::OptionalArg;
8-
use crate::slots::PyIter;
8+
use crate::slots::{IteratorIterable, PyIter};
99
use crate::vm::VirtualMachine;
1010
use crate::{
1111
IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol,
@@ -256,6 +256,7 @@ impl PyAsyncGenASend {
256256
}
257257
}
258258

259+
impl IteratorIterable for PyAsyncGenASend {}
259260
impl PyIter for PyAsyncGenASend {
260261
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
261262
zelf.send(vm.ctx.none(), vm)
@@ -397,6 +398,7 @@ impl PyAsyncGenAThrow {
397398
}
398399
}
399400

401+
impl IteratorIterable for PyAsyncGenAThrow {}
400402
impl PyIter for PyAsyncGenAThrow {
401403
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
402404
zelf.send(vm.ctx.none(), vm)

vm/src/builtins/bytearray.rs

+3-1
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@ use crate::common::lock::{
2121
use crate::function::{ArgIterable, FuncArgs, OptionalArg, OptionalOption};
2222
use crate::sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex};
2323
use crate::slots::{
24-
AsBuffer, Callable, Comparable, Hashable, Iterable, PyComparisonOp, PyIter, Unhashable,
24+
AsBuffer, Callable, Comparable, Hashable, Iterable, IteratorIterable, PyComparisonOp, PyIter,
25+
Unhashable,
2526
};
2627
use crate::utils::Either;
2728
use crate::vm::VirtualMachine;
@@ -747,6 +748,7 @@ impl PyValue for PyByteArrayIterator {
747748

748749
#[pyimpl(with(PyIter))]
749750
impl PyByteArrayIterator {}
751+
impl IteratorIterable for PyByteArrayIterator {}
750752
impl PyIter for PyByteArrayIterator {
751753
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
752754
let pos = zelf.position.fetch_add(1);

vm/src/builtins/bytes.rs

+3-1
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,8 @@ use crate::byteslike::ArgBytesLike;
1313
use crate::common::hash::PyHash;
1414
use crate::function::{ArgIterable, OptionalArg, OptionalOption};
1515
use crate::slots::{
16-
AsBuffer, Callable, Comparable, Hashable, Iterable, PyComparisonOp, PyIter, SlotConstructor,
16+
AsBuffer, Callable, Comparable, Hashable, Iterable, IteratorIterable, PyComparisonOp, PyIter,
17+
SlotConstructor,
1718
};
1819
use crate::utils::Either;
1920
use crate::vm::VirtualMachine;
@@ -601,6 +602,7 @@ impl PyValue for PyBytesIterator {
601602

602603
#[pyimpl(with(PyIter))]
603604
impl PyBytesIterator {}
605+
impl IteratorIterable for PyBytesIterator {}
604606
impl PyIter for PyBytesIterator {
605607
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
606608
let pos = zelf.position.fetch_add(1);

vm/src/builtins/coroutine.rs

+3-1
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use super::pytype::PyTypeRef;
44
use crate::coroutine::{Coro, Variant};
55
use crate::frame::FrameRef;
66
use crate::function::OptionalArg;
7-
use crate::slots::PyIter;
7+
use crate::slots::{IteratorIterable, PyIter};
88
use crate::vm::VirtualMachine;
99
use crate::{IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
1010

@@ -104,6 +104,7 @@ impl PyCoroutine {
104104
}
105105
}
106106

107+
impl IteratorIterable for PyCoroutine {}
107108
impl PyIter for PyCoroutine {
108109
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
109110
zelf.send(vm.ctx.none(), vm)
@@ -141,6 +142,7 @@ impl PyCoroutineWrapper {
141142
}
142143
}
143144

145+
impl IteratorIterable for PyCoroutineWrapper {}
144146
impl PyIter for PyCoroutineWrapper {
145147
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
146148
zelf.send(vm.ctx.none(), vm)

vm/src/builtins/dict.rs

+5-1
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,9 @@ use crate::dictdatatype::{self, DictKey};
1010
use crate::exceptions::PyBaseExceptionRef;
1111
use crate::function::{ArgIterable, FuncArgs, KwArgs, OptionalArg};
1212
use crate::iterator;
13-
use crate::slots::{Comparable, Hashable, Iterable, PyComparisonOp, PyIter, Unhashable};
13+
use crate::slots::{
14+
Comparable, Hashable, Iterable, IteratorIterable, PyComparisonOp, PyIter, Unhashable,
15+
};
1416
use crate::vm::{ReprGuard, VirtualMachine};
1517
use crate::{
1618
IdProtocol, IntoPyObject, ItemProtocol, PyArithmaticValue::*, PyAttributes, PyClassDef,
@@ -731,6 +733,7 @@ macro_rules! dict_iterator {
731733
}
732734
}
733735

736+
impl IteratorIterable for $iter_name {}
734737
impl PyIter for $iter_name {
735738
#[allow(clippy::redundant_closure_call)]
736739
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
@@ -791,6 +794,7 @@ macro_rules! dict_iterator {
791794
}
792795
}
793796

797+
impl IteratorIterable for $reverse_iter_name {}
794798
impl PyIter for $reverse_iter_name {
795799
#[allow(clippy::redundant_closure_call)]
796800
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {

vm/src/builtins/enumerate.rs

+3-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ use super::iter::{
1111
};
1212
use super::pytype::PyTypeRef;
1313
use crate::function::OptionalArg;
14-
use crate::slots::{PyIter, SlotConstructor};
14+
use crate::slots::{IteratorIterable, PyIter, SlotConstructor};
1515
use crate::vm::VirtualMachine;
1616
use crate::{iterator, ItemProtocol, TypeProtocol};
1717
use crate::{IntoPyObject, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
@@ -56,6 +56,7 @@ impl SlotConstructor for PyEnumerate {
5656
#[pyimpl(with(PyIter, SlotConstructor), flags(BASETYPE))]
5757
impl PyEnumerate {}
5858

59+
impl IteratorIterable for PyEnumerate {}
5960
impl PyIter for PyEnumerate {
6061
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
6162
let next_obj = iterator::call_next(vm, &zelf.iterator)?;
@@ -137,6 +138,7 @@ impl PyReverseSequenceIterator {
137138
}
138139
}
139140

141+
impl IteratorIterable for PyReverseSequenceIterator {}
140142
impl PyIter for PyReverseSequenceIterator {
141143
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
142144
if let Exhausted = zelf.status.load() {

vm/src/builtins/filter.rs

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use super::pytype::PyTypeRef;
22
use crate::iterator;
3-
use crate::slots::{PyIter, SlotConstructor};
3+
use crate::slots::{IteratorIterable, PyIter, SlotConstructor};
44
use crate::vm::VirtualMachine;
55
use crate::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
66

@@ -50,6 +50,7 @@ impl SlotConstructor for PyFilter {
5050
#[pyimpl(with(PyIter, SlotConstructor), flags(BASETYPE))]
5151
impl PyFilter {}
5252

53+
impl IteratorIterable for PyFilter {}
5354
impl PyIter for PyFilter {
5455
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
5556
let predicate = &zelf.predicate;

vm/src/builtins/generator.rs

+2-1
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use super::pytype::PyTypeRef;
88
use crate::coroutine::{Coro, Variant};
99
use crate::frame::FrameRef;
1010
use crate::function::OptionalArg;
11-
use crate::slots::PyIter;
11+
use crate::slots::{IteratorIterable, PyIter};
1212
use crate::vm::VirtualMachine;
1313
use crate::{IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
1414

@@ -95,6 +95,7 @@ impl PyGenerator {
9595
}
9696
}
9797

98+
impl IteratorIterable for PyGenerator {}
9899
impl PyIter for PyGenerator {
99100
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
100101
zelf.send(vm.ctx.none(), vm)

vm/src/builtins/iter.rs

+3-1
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use crossbeam_utils::atomic::AtomicCell;
66

77
use super::pytype::PyTypeRef;
88
use super::{int, PyInt};
9-
use crate::slots::PyIter;
9+
use crate::slots::{IteratorIterable, PyIter};
1010
use crate::vm::VirtualMachine;
1111
use crate::{
1212
function::ArgCallable, ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult,
@@ -92,6 +92,7 @@ impl PySequenceIterator {
9292
}
9393
}
9494

95+
impl IteratorIterable for PySequenceIterator {}
9596
impl PyIter for PySequenceIterator {
9697
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
9798
if let IterStatus::Exhausted = zelf.status.load() {
@@ -134,6 +135,7 @@ impl PyCallableIterator {
134135
}
135136
}
136137

138+
impl IteratorIterable for PyCallableIterator {}
137139
impl PyIter for PyCallableIterator {
138140
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
139141
if let IterStatus::Exhausted = zelf.status.load() {

vm/src/builtins/list.rs

+5-1
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,9 @@ use crate::common::lock::{
2020
use crate::function::{ArgIterable, FuncArgs, OptionalArg};
2121
use crate::sequence::{self, SimpleSeq};
2222
use crate::sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex};
23-
use crate::slots::{Comparable, Hashable, Iterable, PyComparisonOp, PyIter, Unhashable};
23+
use crate::slots::{
24+
Comparable, Hashable, Iterable, IteratorIterable, PyComparisonOp, PyIter, Unhashable,
25+
};
2426
use crate::utils::Either;
2527
use crate::vm::{ReprGuard, VirtualMachine};
2628
use crate::{
@@ -526,6 +528,7 @@ impl PyListIterator {
526528
}
527529
}
528530

531+
impl IteratorIterable for PyListIterator {}
529532
impl PyIter for PyListIterator {
530533
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
531534
if let Exhausted = zelf.status.load() {
@@ -599,6 +602,7 @@ impl PyListReverseIterator {
599602
}
600603
}
601604

605+
impl IteratorIterable for PyListReverseIterator {}
602606
impl PyIter for PyListReverseIterator {
603607
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
604608
if let Exhausted = zelf.status.load() {

vm/src/builtins/map.rs

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use super::pytype::PyTypeRef;
22
use crate::function::PosArgs;
33
use crate::iterator;
4-
use crate::slots::{PyIter, SlotConstructor};
4+
use crate::slots::{IteratorIterable, PyIter, SlotConstructor};
55
use crate::vm::VirtualMachine;
66
use crate::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
77

@@ -50,6 +50,7 @@ impl PyMap {
5050
}
5151
}
5252

53+
impl IteratorIterable for PyMap {}
5354
impl PyIter for PyMap {
5455
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
5556
let next_objs = zelf

vm/src/builtins/pystr.rs

+4-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,9 @@ use crate::exceptions::IntoPyException;
1111
use crate::format::{FormatSpec, FormatString, FromTemplate};
1212
use crate::function::{ArgIterable, FuncArgs, OptionalArg, OptionalOption};
1313
use crate::sliceable::PySliceableSequence;
14-
use crate::slots::{Comparable, Hashable, Iterable, PyComparisonOp, PyIter, SlotConstructor};
14+
use crate::slots::{
15+
Comparable, Hashable, Iterable, IteratorIterable, PyComparisonOp, PyIter, SlotConstructor,
16+
};
1517
use crate::utils::Either;
1618
use crate::VirtualMachine;
1719
use crate::{
@@ -227,6 +229,7 @@ impl PyStrIterator {
227229
}
228230
}
229231

232+
impl IteratorIterable for PyStrIterator {}
230233
impl PyIter for PyStrIterator {
231234
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
232235
if let Exhausted = zelf.status.load() {

vm/src/builtins/range.rs

+3-1
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use super::pytype::PyTypeRef;
99
use super::slice::{PySlice, PySliceRef};
1010
use crate::common::hash::PyHash;
1111
use crate::function::{FuncArgs, OptionalArg};
12-
use crate::slots::{Comparable, Hashable, Iterable, PyComparisonOp, PyIter};
12+
use crate::slots::{Comparable, Hashable, Iterable, IteratorIterable, PyComparisonOp, PyIter};
1313
use crate::vm::VirtualMachine;
1414
use crate::{
1515
iterator, IdProtocol, IntoPyRef, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue,
@@ -516,6 +516,7 @@ impl PyLongRangeIterator {
516516
}
517517
}
518518

519+
impl IteratorIterable for PyLongRangeIterator {}
519520
impl PyIter for PyLongRangeIterator {
520521
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
521522
// TODO: In pathological case (index == usize::MAX) this can wrap around
@@ -585,6 +586,7 @@ impl PyRangeIterator {
585586
}
586587
}
587588

589+
impl IteratorIterable for PyRangeIterator {}
588590
impl PyIter for PyRangeIterator {
589591
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
590592
// TODO: In pathological case (index == usize::MAX) this can wrap around

vm/src/builtins/set.rs

+3-1
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,8 @@ use crate::dictdatatype;
88
use crate::dictdatatype::DictSize;
99
use crate::function::{ArgIterable, FuncArgs, OptionalArg, PosArgs};
1010
use crate::slots::{
11-
Comparable, Hashable, Iterable, PyComparisonOp, PyIter, SlotConstructor, Unhashable,
11+
Comparable, Hashable, Iterable, IteratorIterable, PyComparisonOp, PyIter, SlotConstructor,
12+
Unhashable,
1213
};
1314
use crate::vm::{ReprGuard, VirtualMachine};
1415
use crate::{
@@ -864,6 +865,7 @@ impl PySetIterator {
864865
}
865866
}
866867

868+
impl IteratorIterable for PySetIterator {}
867869
impl PyIter for PySetIterator {
868870
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
869871
match zelf.status.load() {

vm/src/builtins/tuple.rs

+4-1
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,9 @@ use crate::common::hash::PyHash;
1515
use crate::function::OptionalArg;
1616
use crate::sequence::{self, SimpleSeq};
1717
use crate::sliceable::PySliceableSequence;
18-
use crate::slots::{Comparable, Hashable, Iterable, PyComparisonOp, PyIter, SlotConstructor};
18+
use crate::slots::{
19+
Comparable, Hashable, Iterable, IteratorIterable, PyComparisonOp, PyIter, SlotConstructor,
20+
};
1921
use crate::utils::Either;
2022
use crate::vm::{ReprGuard, VirtualMachine};
2123
use crate::{
@@ -378,6 +380,7 @@ impl PyTupleIterator {
378380
}
379381
}
380382

383+
impl IteratorIterable for PyTupleIterator {}
381384
impl PyIter for PyTupleIterator {
382385
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
383386
if let Exhausted = zelf.status.load() {

vm/src/builtins/zip.rs

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use super::pytype::PyTypeRef;
22
use crate::function::PosArgs;
33
use crate::iterator;
4-
use crate::slots::{PyIter, SlotConstructor};
4+
use crate::slots::{IteratorIterable, PyIter, SlotConstructor};
55
use crate::vm::VirtualMachine;
66
use crate::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
77

@@ -32,6 +32,7 @@ impl SlotConstructor for PyZip {
3232
#[pyimpl(with(PyIter, SlotConstructor), flags(BASETYPE))]
3333
impl PyZip {}
3434

35+
impl IteratorIterable for PyZip {}
3536
impl PyIter for PyZip {
3637
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
3738
if zelf.iterators.is_empty() {

vm/src/slots.rs

+9-7
Original file line numberDiff line numberDiff line change
@@ -533,6 +533,7 @@ pub trait AsBuffer: PyValue {
533533
#[pyimpl]
534534
pub trait Iterable: PyValue {
535535
#[pyslot]
536+
#[pymethod(name = "__iter__")]
536537
fn tp_iter(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult {
537538
if let Ok(zelf) = zelf.downcast() {
538539
Self::iter(zelf, vm)
@@ -541,12 +542,11 @@ pub trait Iterable: PyValue {
541542
}
542543
}
543544

544-
#[pymethod(magic)]
545545
fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult;
546546
}
547547

548548
#[pyimpl(with(Iterable))]
549-
pub trait PyIter: PyValue {
549+
pub trait PyIter: PyValue + Iterable {
550550
#[pyslot]
551551
fn tp_iternext(zelf: &PyObjectRef, vm: &VirtualMachine) -> PyResult {
552552
if let Some(zelf) = zelf.downcast_ref() {
@@ -559,19 +559,21 @@ pub trait PyIter: PyValue {
559559
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult;
560560

561561
#[pymethod]
562-
fn __next__(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
563-
Self::next(&zelf, vm)
562+
fn __next__(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult {
563+
Self::tp_iternext(&zelf, vm)
564564
}
565565
}
566566

567+
pub trait IteratorIterable: PyValue {}
568+
567569
impl<T> Iterable for T
568570
where
569-
T: PyIter,
571+
T: IteratorIterable,
570572
{
571573
fn tp_iter(zelf: PyObjectRef, _vm: &VirtualMachine) -> PyResult {
572574
Ok(zelf)
573575
}
574-
fn iter(zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyResult {
575-
Ok(zelf.into_object())
576+
fn iter(_zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyResult {
577+
unreachable!("tp_iter is implemented");
576578
}
577579
}

0 commit comments

Comments
 (0)