Skip to content

Commit 6d7eae9

Browse files
coolreader18coolreader18
authored
Add iter, iternext slots & Allow deleting slot magic methods
1 parent 64b951e commit 6d7eae9

34 files changed

+848
-918
lines changed

‎Lib/test/test_types.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1456,6 +1456,7 @@ def foo():
14561456
self.assertIs(foo(), coro)
14571457
self.assertIs(foo().__await__(), coro)
14581458

1459+
@unittest.skip("TODO: RUSTPYTHON, unittest.mock")
14591460
def test_duck_gen(self):
14601461
class GenLike:
14611462
def send(self): pass

‎derive/src/pyclass.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,12 @@ pub(crate) fn impl_pyimpl(
9292
let mut context = ImplContext::default();
9393
extract_items_into_context(&mut context, trai.items.iter_mut());
9494

95+
let ExtractedImplAttrs {
96+
with_impl,
97+
with_slots,
98+
..
99+
} = extract_impl_attrs(attr)?;
100+
95101
let getset_impl = &context.getset_items;
96102
let extend_impl = &context.impl_extend_items;
97103
let slots_impl = &context.extend_slots_items;
@@ -104,11 +110,13 @@ pub(crate) fn impl_pyimpl(
104110
) {
105111
#getset_impl
106112
#extend_impl
113+
#with_impl
107114
#(#class_extensions)*
108115
}
109116
},
110117
parse_quote! {
111118
fn __extend_slots(slots: &mut ::rustpython_vm::slots::PyTypeSlots) {
119+
#with_slots
112120
#slots_impl
113121
}
114122
},

‎vm/src/builtins/asyncgenerator.rs‎

Lines changed: 15 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use crate::function::OptionalArg;
77
use crate::pyobject::{
88
PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol,
99
};
10+
use crate::slots::PyIter;
1011
use crate::vm::VirtualMachine;
1112

1213
use crossbeam_utils::atomic::AtomicCell;
@@ -172,22 +173,13 @@ impl PyValue for PyAsyncGenASend {
172173
}
173174
}
174175

175-
#[pyimpl]
176+
#[pyimpl(with(PyIter))]
176177
impl PyAsyncGenASend {
177-
#[pymethod(name = "__iter__")]
178-
fn iter(zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyRef<Self> {
179-
zelf
180-
}
181178
#[pymethod(name = "__await__")]
182179
fn r#await(zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyRef<Self> {
183180
zelf
184181
}
185182

186-
#[pymethod(name = "__next__")]
187-
fn next(&self, vm: &VirtualMachine) -> PyResult {
188-
self.send(vm.ctx.none(), vm)
189-
}
190-
191183
#[pymethod]
192184
fn send(&self, val: PyObjectRef, vm: &VirtualMachine) -> PyResult {
193185
let val = match self.state.load() {
@@ -253,6 +245,12 @@ impl PyAsyncGenASend {
253245
}
254246
}
255247

248+
impl PyIter for PyAsyncGenASend {
249+
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
250+
zelf.send(vm.ctx.none(), vm)
251+
}
252+
}
253+
256254
#[pyclass(module = false, name = "async_generator_athrow")]
257255
#[derive(Debug)]
258256
pub(crate) struct PyAsyncGenAThrow {
@@ -268,22 +266,13 @@ impl PyValue for PyAsyncGenAThrow {
268266
}
269267
}
270268

271-
#[pyimpl]
269+
#[pyimpl(with(PyIter))]
272270
impl PyAsyncGenAThrow {
273-
#[pymethod(name = "__iter__")]
274-
fn iter(zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyRef<Self> {
275-
zelf
276-
}
277271
#[pymethod(name = "__await__")]
278272
fn r#await(zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyRef<Self> {
279273
zelf
280274
}
281275

282-
#[pymethod(name = "__next__")]
283-
fn next(&self, vm: &VirtualMachine) -> PyResult {
284-
self.send(vm.ctx.none(), vm)
285-
}
286-
287276
#[pymethod]
288277
fn send(&self, val: PyObjectRef, vm: &VirtualMachine) -> PyResult {
289278
match self.state.load() {
@@ -397,6 +386,12 @@ impl PyAsyncGenAThrow {
397386
}
398387
}
399388

389+
impl PyIter for PyAsyncGenAThrow {
390+
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
391+
zelf.send(vm.ctx.none(), vm)
392+
}
393+
}
394+
400395
pub fn init(ctx: &PyContext) {
401396
PyAsyncGen::extend_class(ctx, &ctx.types.async_generator);
402397
PyAsyncGenASend::extend_class(ctx, &ctx.types.async_generator_asend);

‎vm/src/builtins/bytearray.rs‎

Lines changed: 21 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,9 @@ use crate::pyobject::{
2121
PyIterable, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol,
2222
};
2323
use crate::sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex};
24-
use crate::slots::{BufferProtocol, Comparable, Hashable, PyComparisonOp, Unhashable};
24+
use crate::slots::{
25+
BufferProtocol, Comparable, Hashable, Iterable, PyComparisonOp, PyIter, Unhashable,
26+
};
2527
use crate::vm::VirtualMachine;
2628
use bstr::ByteSlice;
2729
use crossbeam_utils::atomic::AtomicCell;
@@ -97,7 +99,7 @@ pub(crate) fn init(context: &PyContext) {
9799
PyByteArrayIterator::extend_class(context, &context.types.bytearray_iterator_type);
98100
}
99101

100-
#[pyimpl(flags(BASETYPE), with(Hashable, Comparable, BufferProtocol))]
102+
#[pyimpl(flags(BASETYPE), with(Hashable, Comparable, BufferProtocol, Iterable))]
101103
impl PyByteArray {
102104
#[pyslot]
103105
fn tp_new(
@@ -123,14 +125,6 @@ impl PyByteArray {
123125
size_of::<Self>() + self.borrow_value().len() * size_of::<u8>()
124126
}
125127

126-
#[pymethod(name = "__iter__")]
127-
fn iter(zelf: PyRef<Self>) -> PyByteArrayIterator {
128-
PyByteArrayIterator {
129-
position: AtomicCell::new(0),
130-
bytearray: zelf,
131-
}
132-
}
133-
134128
#[pymethod(name = "__add__")]
135129
fn add(&self, other: PyBytesLike, vm: &VirtualMachine) -> PyObjectRef {
136130
vm.ctx
@@ -718,6 +712,16 @@ impl<'a> ResizeGuard<'a> for PyByteArray {
718712

719713
impl Unhashable for PyByteArray {}
720714

715+
impl Iterable for PyByteArray {
716+
fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
717+
Ok(PyByteArrayIterator {
718+
position: AtomicCell::new(0),
719+
bytearray: zelf,
720+
}
721+
.into_object(vm))
722+
}
723+
}
724+
721725
// fn set_value(obj: &PyObjectRef, value: Vec<u8>) {
722726
// obj.borrow_mut().kind = PyObjectPayload::Bytes { value };
723727
// }
@@ -735,20 +739,15 @@ impl PyValue for PyByteArrayIterator {
735739
}
736740
}
737741

738-
#[pyimpl]
739-
impl PyByteArrayIterator {
740-
#[pymethod(name = "__next__")]
741-
fn next(&self, vm: &VirtualMachine) -> PyResult<u8> {
742-
let pos = self.position.fetch_add(1);
743-
if let Some(&ret) = self.bytearray.borrow_value().elements.get(pos) {
744-
Ok(ret)
742+
#[pyimpl(with(PyIter))]
743+
impl PyByteArrayIterator {}
744+
impl PyIter for PyByteArrayIterator {
745+
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
746+
let pos = zelf.position.fetch_add(1);
747+
if let Some(&ret) = zelf.bytearray.borrow_value().elements.get(pos) {
748+
Ok(ret.into_pyobject(vm))
745749
} else {
746750
Err(vm.new_stop_iteration())
747751
}
748752
}
749-
750-
#[pymethod(name = "__iter__")]
751-
fn iter(zelf: PyRef<Self>) -> PyRef<Self> {
752-
zelf
753-
}
754753
}

‎vm/src/builtins/bytes.rs‎

Lines changed: 19 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ use crate::pyobject::{
1919
BorrowValue, Either, IntoPyObject, PyClassImpl, PyComparisonValue, PyContext, PyIterable,
2020
PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol,
2121
};
22-
use crate::slots::{BufferProtocol, Comparable, Hashable, PyComparisonOp};
22+
use crate::slots::{BufferProtocol, Comparable, Hashable, Iterable, PyComparisonOp, PyIter};
2323
use crate::vm::VirtualMachine;
2424
use crate::{
2525
anystr::{self, AnyStr},
@@ -96,7 +96,7 @@ pub(crate) fn init(context: &PyContext) {
9696
PyBytesIterator::extend_class(context, &context.types.bytes_iterator_type);
9797
}
9898

99-
#[pyimpl(flags(BASETYPE), with(Hashable, Comparable, BufferProtocol))]
99+
#[pyimpl(flags(BASETYPE), with(Hashable, Comparable, BufferProtocol, Iterable))]
100100
impl PyBytes {
101101
#[pyslot]
102102
fn tp_new(
@@ -117,14 +117,6 @@ impl PyBytes {
117117
self.inner.len()
118118
}
119119

120-
#[pymethod(name = "__iter__")]
121-
fn iter(zelf: PyRef<Self>) -> PyBytesIterator {
122-
PyBytesIterator {
123-
position: AtomicCell::new(0),
124-
bytes: zelf,
125-
}
126-
}
127-
128120
#[pymethod(name = "__sizeof__")]
129121
fn sizeof(&self) -> PyResult<usize> {
130122
Ok(size_of::<Self>() + self.inner.elements.len() * size_of::<u8>())
@@ -545,6 +537,16 @@ impl Comparable for PyBytes {
545537
}
546538
}
547539

540+
impl Iterable for PyBytes {
541+
fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
542+
Ok(PyBytesIterator {
543+
position: AtomicCell::new(0),
544+
bytes: zelf,
545+
}
546+
.into_object(vm))
547+
}
548+
}
549+
548550
#[pyclass(module = false, name = "bytes_iterator")]
549551
#[derive(Debug)]
550552
pub struct PyBytesIterator {
@@ -558,22 +560,17 @@ impl PyValue for PyBytesIterator {
558560
}
559561
}
560562

561-
#[pyimpl]
562-
impl PyBytesIterator {
563-
#[pymethod(name = "__next__")]
564-
fn next(&self, vm: &VirtualMachine) -> PyResult<u8> {
565-
let pos = self.position.fetch_add(1);
566-
if let Some(&ret) = self.bytes.borrow_value().get(pos) {
567-
Ok(ret)
563+
#[pyimpl(with(PyIter))]
564+
impl PyBytesIterator {}
565+
impl PyIter for PyBytesIterator {
566+
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
567+
let pos = zelf.position.fetch_add(1);
568+
if let Some(&ret) = zelf.bytes.borrow_value().get(pos) {
569+
Ok(vm.ctx.new_int(ret))
568570
} else {
569571
Err(vm.new_stop_iteration())
570572
}
571573
}
572-
573-
#[pymethod(name = "__iter__")]
574-
fn iter(zelf: PyRef<Self>) -> PyRef<Self> {
575-
zelf
576-
}
577574
}
578575

579576
impl TryFromObject for PyBytes {

‎vm/src/builtins/coroutine.rs‎

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use crate::coroutine::{Coro, Variant};
55
use crate::frame::FrameRef;
66
use crate::function::OptionalArg;
77
use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
8+
use crate::slots::PyIter;
89
use crate::vm::VirtualMachine;
910

1011
type PyCoroutineRef = PyRef<PyCoroutine>;
@@ -21,7 +22,7 @@ impl PyValue for PyCoroutine {
2122
}
2223
}
2324

24-
#[pyimpl]
25+
#[pyimpl(with(PyIter))]
2526
impl PyCoroutine {
2627
pub fn as_coro(&self) -> &Coro {
2728
&self.inner
@@ -95,6 +96,12 @@ impl PyCoroutine {
9596
}
9697
}
9798

99+
impl PyIter for PyCoroutine {
100+
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
101+
zelf.send(vm.ctx.none(), vm)
102+
}
103+
}
104+
98105
#[pyclass(module = false, name = "coroutine_wrapper")]
99106
#[derive(Debug)]
100107
pub struct PyCoroutineWrapper {
@@ -109,16 +116,6 @@ impl PyValue for PyCoroutineWrapper {
109116

110117
#[pyimpl]
111118
impl PyCoroutineWrapper {
112-
#[pymethod(name = "__iter__")]
113-
fn iter(zelf: PyRef<Self>) -> PyRef<Self> {
114-
zelf
115-
}
116-
117-
#[pymethod(name = "__next__")]
118-
fn next(&self, vm: &VirtualMachine) -> PyResult {
119-
self.coro.send(vm.ctx.none(), vm)
120-
}
121-
122119
#[pymethod]
123120
fn send(&self, val: PyObjectRef, vm: &VirtualMachine) -> PyResult {
124121
self.coro.send(val, vm)

0 commit comments

Comments
 (0)