Skip to content

Commit 3eec37b

Browse files
youknowoneyouknowone
committed
use PyIterReturn for coroutine internal
1 parent 54f08cc commit 3eec37b

11 files changed

Lines changed: 161 additions & 163 deletions

File tree

‎vm/src/builtins/asyncgenerator.rs‎

Lines changed: 34 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use super::{PyCode, PyStrRef, PyTypeRef};
22
use crate::{
33
builtins::PyBaseExceptionRef,
4-
coroutine::{Coro, Variant},
4+
coroutine::Coro,
55
frame::FrameRef,
66
function::OptionalArg,
77
protocol::PyIterReturn,
@@ -34,7 +34,7 @@ impl PyAsyncGen {
3434

3535
pub fn new(frame: FrameRef, name: PyStrRef) -> Self {
3636
PyAsyncGen {
37-
inner: Coro::new(frame, Variant::AsyncGen, name),
37+
inner: Coro::new(frame, name),
3838
running_async: AtomicCell::new(false),
3939
}
4040
}
@@ -50,8 +50,8 @@ impl PyAsyncGen {
5050
}
5151

5252
#[pymethod(magic)]
53-
fn repr(zelf: PyRef<Self>) -> String {
54-
zelf.inner.repr(zelf.get_id())
53+
fn repr(zelf: PyRef<Self>, vm: &VirtualMachine) -> String {
54+
zelf.inner.repr(zelf.as_object(), zelf.get_id(), vm)
5555
}
5656

5757
#[pymethod(magic)]
@@ -138,17 +138,20 @@ impl PyValue for PyAsyncGenWrappedValue {
138138
impl PyAsyncGenWrappedValue {}
139139

140140
impl PyAsyncGenWrappedValue {
141-
fn unbox(ag: &PyAsyncGen, val: PyResult, vm: &VirtualMachine) -> PyResult {
142-
if let Err(ref e) = val {
143-
if e.isinstance(&vm.ctx.exceptions.stop_async_iteration)
144-
|| e.isinstance(&vm.ctx.exceptions.generator_exit)
145-
{
146-
ag.inner.closed.store(true);
147-
}
141+
fn unbox(ag: &PyAsyncGen, val: PyResult<PyIterReturn>, vm: &VirtualMachine) -> PyResult {
142+
let (closed, async_done) = match &val {
143+
Ok(PyIterReturn::StopIteration(_)) => (true, true),
144+
Err(e) if e.isinstance(&vm.ctx.exceptions.generator_exit) => (true, true),
145+
Err(_) => (false, true),
146+
_ => (false, false),
147+
};
148+
if closed {
149+
ag.inner.closed.store(true);
150+
}
151+
if async_done {
148152
ag.running_async.store(false);
149153
}
150-
let val = val?;
151-
154+
let val = val?.into_async_pyresult(vm)?;
152155
match_class!(match val {
153156
val @ Self => {
154157
ag.running_async.store(false);
@@ -214,7 +217,7 @@ impl PyAsyncGenASend {
214217
}
215218
}
216219
};
217-
let res = self.ag.inner.send(val, vm);
220+
let res = self.ag.inner.send(self.ag.as_object(), val, vm);
218221
let res = PyAsyncGenWrappedValue::unbox(&self.ag, res, vm);
219222
if res.is_err() {
220223
self.close();
@@ -237,6 +240,7 @@ impl PyAsyncGenASend {
237240
}
238241

239242
let res = self.ag.inner.throw(
243+
self.ag.as_object(),
240244
exc_type,
241245
exc_val.unwrap_or_none(vm),
242246
exc_tb.unwrap_or_none(vm),
@@ -258,8 +262,7 @@ impl PyAsyncGenASend {
258262
impl IteratorIterable for PyAsyncGenASend {}
259263
impl SlotIterator for PyAsyncGenASend {
260264
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
261-
// TODO: Fix zelf.send to return PyIterReturn
262-
PyIterReturn::from_result(zelf.send(vm.ctx.none(), vm), vm)
265+
PyIterReturn::from_pyresult(zelf.send(vm.ctx.none(), vm), vm)
263266
}
264267
}
265268

@@ -315,27 +318,28 @@ impl PyAsyncGenAThrow {
315318
self.ag.running_async.store(true);
316319

317320
let (ty, val, tb) = self.value.clone();
318-
let ret = self.ag.inner.throw(ty, val, tb, vm);
321+
let ret = self.ag.inner.throw(self.ag.as_object(), ty, val, tb, vm);
319322
let ret = if self.aclose {
320323
if self.ignored_close(&ret) {
321324
Err(self.yield_close(vm))
322325
} else {
323-
ret
326+
ret.and_then(|o| o.into_async_pyresult(vm))
324327
}
325328
} else {
326329
PyAsyncGenWrappedValue::unbox(&self.ag, ret, vm)
327330
};
328331
ret.map_err(|e| self.check_error(e, vm))
329332
}
330333
AwaitableState::Iter => {
331-
let ret = self.ag.inner.send(val, vm);
334+
let ret = self.ag.inner.send(self.ag.as_object(), val, vm);
332335
if self.aclose {
333336
match ret {
334-
Ok(v) if v.payload_is::<PyAsyncGenWrappedValue>() => {
337+
Ok(PyIterReturn::Return(v)) if v.payload_is::<PyAsyncGenWrappedValue>() => {
335338
Err(self.yield_close(vm))
336339
}
337-
Ok(v) => Ok(v),
338-
Err(e) => Err(self.check_error(e, vm)),
340+
other => other
341+
.and_then(|o| o.into_async_pyresult(vm))
342+
.map_err(|e| self.check_error(e, vm)),
339343
}
340344
} else {
341345
PyAsyncGenWrappedValue::unbox(&self.ag, ret, vm)
@@ -353,6 +357,7 @@ impl PyAsyncGenAThrow {
353357
vm: &VirtualMachine,
354358
) -> PyResult {
355359
let ret = self.ag.inner.throw(
360+
self.ag.as_object(),
356361
exc_type,
357362
exc_val.unwrap_or_none(vm),
358363
exc_tb.unwrap_or_none(vm),
@@ -362,7 +367,7 @@ impl PyAsyncGenAThrow {
362367
if self.ignored_close(&ret) {
363368
Err(self.yield_close(vm))
364369
} else {
365-
ret
370+
ret.and_then(|o| o.into_async_pyresult(vm))
366371
}
367372
} else {
368373
PyAsyncGenWrappedValue::unbox(&self.ag, ret, vm)
@@ -375,9 +380,11 @@ impl PyAsyncGenAThrow {
375380
self.state.store(AwaitableState::Closed);
376381
}
377382

378-
fn ignored_close(&self, res: &PyResult) -> bool {
379-
res.as_ref()
380-
.map_or(false, |v| v.payload_is::<PyAsyncGenWrappedValue>())
383+
fn ignored_close(&self, res: &PyResult<PyIterReturn>) -> bool {
384+
res.as_ref().map_or(false, |v| match v {
385+
PyIterReturn::Return(obj) => obj.payload_is::<PyAsyncGenWrappedValue>(),
386+
PyIterReturn::StopIteration(_) => false,
387+
})
381388
}
382389
fn yield_close(&self, vm: &VirtualMachine) -> PyBaseExceptionRef {
383390
self.ag.running_async.store(false);
@@ -401,8 +408,7 @@ impl PyAsyncGenAThrow {
401408
impl IteratorIterable for PyAsyncGenAThrow {}
402409
impl SlotIterator for PyAsyncGenAThrow {
403410
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
404-
// TODO: Fix zelf.send to return PyIterReturn
405-
PyIterReturn::from_result(zelf.send(vm.ctx.none(), vm), vm)
411+
PyIterReturn::from_pyresult(zelf.send(vm.ctx.none(), vm), vm)
406412
}
407413
}
408414

‎vm/src/builtins/coroutine.rs‎

Lines changed: 18 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use super::{PyCode, PyStrRef, PyTypeRef};
22
use crate::{
3-
coroutine::{Coro, Variant},
3+
coroutine::Coro,
44
frame::FrameRef,
55
function::OptionalArg,
66
protocol::PyIterReturn,
@@ -29,7 +29,7 @@ impl PyCoroutine {
2929

3030
pub fn new(frame: FrameRef, name: PyStrRef) -> Self {
3131
PyCoroutine {
32-
inner: Coro::new(frame, Variant::Coroutine, name),
32+
inner: Coro::new(frame, name),
3333
}
3434
}
3535

@@ -44,24 +44,25 @@ impl PyCoroutine {
4444
}
4545

4646
#[pymethod(magic)]
47-
fn repr(zelf: PyRef<Self>) -> String {
48-
zelf.inner.repr(zelf.get_id())
47+
fn repr(zelf: PyRef<Self>, vm: &VirtualMachine) -> String {
48+
zelf.inner.repr(zelf.as_object(), zelf.get_id(), vm)
4949
}
5050

5151
#[pymethod]
52-
fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult {
53-
self.inner.send(value, vm)
52+
fn send(zelf: PyRef<Self>, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
53+
zelf.inner.send(zelf.as_object(), value, vm)
5454
}
5555

5656
#[pymethod]
5757
fn throw(
58-
&self,
58+
zelf: PyRef<Self>,
5959
exc_type: PyObjectRef,
6060
exc_val: OptionalArg,
6161
exc_tb: OptionalArg,
6262
vm: &VirtualMachine,
63-
) -> PyResult {
64-
self.inner.throw(
63+
) -> PyResult<PyIterReturn> {
64+
zelf.inner.throw(
65+
zelf.as_object(),
6566
exc_type,
6667
exc_val.unwrap_or_none(vm),
6768
exc_tb.unwrap_or_none(vm),
@@ -70,8 +71,8 @@ impl PyCoroutine {
7071
}
7172

7273
#[pymethod]
73-
fn close(&self, vm: &VirtualMachine) -> PyResult<()> {
74-
self.inner.close(vm)
74+
fn close(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<()> {
75+
zelf.inner.close(zelf.as_object(), vm)
7576
}
7677

7778
#[pymethod(name = "__await__")]
@@ -106,8 +107,7 @@ impl PyCoroutine {
106107
impl IteratorIterable for PyCoroutine {}
107108
impl SlotIterator for PyCoroutine {
108109
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
109-
// TODO: Fix zelf.send to return PyIterReturn
110-
PyIterReturn::from_result(zelf.send(vm.ctx.none(), vm), vm)
110+
Self::send(zelf.clone(), vm.ctx.none(), vm)
111111
}
112112
}
113113

@@ -127,8 +127,8 @@ impl PyValue for PyCoroutineWrapper {
127127
#[pyimpl(with(SlotIterator))]
128128
impl PyCoroutineWrapper {
129129
#[pymethod]
130-
fn send(&self, val: PyObjectRef, vm: &VirtualMachine) -> PyResult {
131-
self.coro.send(val, vm)
130+
fn send(zelf: PyRef<Self>, val: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
131+
PyCoroutine::send(zelf.coro.clone(), val, vm)
132132
}
133133

134134
#[pymethod]
@@ -138,16 +138,15 @@ impl PyCoroutineWrapper {
138138
exc_val: OptionalArg,
139139
exc_tb: OptionalArg,
140140
vm: &VirtualMachine,
141-
) -> PyResult {
142-
self.coro.throw(exc_type, exc_val, exc_tb, vm)
141+
) -> PyResult<PyIterReturn> {
142+
PyCoroutine::throw(self.coro.clone(), exc_type, exc_val, exc_tb, vm)
143143
}
144144
}
145145

146146
impl IteratorIterable for PyCoroutineWrapper {}
147147
impl SlotIterator for PyCoroutineWrapper {
148148
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
149-
// TODO: Fix zelf.send to return PyIterReturn
150-
PyIterReturn::from_result(zelf.send(vm.ctx.none(), vm), vm)
149+
Self::send(zelf.clone(), vm.ctx.none(), vm)
151150
}
152151
}
153152

‎vm/src/builtins/filter.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ impl SlotIterator for PyFilter {
5151
} else {
5252
// the predicate itself can raise StopIteration which does stop the filter
5353
// iteration
54-
match PyIterReturn::from_result(vm.invoke(predicate, vec![next_obj.clone()]), vm)? {
54+
match PyIterReturn::from_pyresult(vm.invoke(predicate, vec![next_obj.clone()]), vm)? {
5555
PyIterReturn::Return(obj) => obj,
5656
PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)),
5757
}

‎vm/src/builtins/generator.rs‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
use super::{PyCode, PyStrRef, PyTypeRef};
66
use crate::{
7-
coroutine::{Coro, Variant},
7+
coroutine::Coro,
88
frame::FrameRef,
99
function::OptionalArg,
1010
protocol::PyIterReturn,
@@ -32,7 +32,7 @@ impl PyGenerator {
3232

3333
pub fn new(frame: FrameRef, name: PyStrRef) -> Self {
3434
PyGenerator {
35-
inner: Coro::new(frame, Variant::Gen, name),
35+
inner: Coro::new(frame, name),
3636
}
3737
}
3838

@@ -47,24 +47,25 @@ impl PyGenerator {
4747
}
4848

4949
#[pymethod(magic)]
50-
fn repr(zelf: PyRef<Self>) -> String {
51-
zelf.inner.repr(zelf.get_id())
50+
fn repr(zelf: PyRef<Self>, vm: &VirtualMachine) -> String {
51+
zelf.inner.repr(zelf.as_object(), zelf.get_id(), vm)
5252
}
5353

5454
#[pymethod]
55-
fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult {
56-
self.inner.send(value, vm)
55+
fn send(zelf: PyRef<Self>, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
56+
zelf.inner.send(zelf.as_object(), value, vm)
5757
}
5858

5959
#[pymethod]
6060
fn throw(
61-
&self,
61+
zelf: PyRef<Self>,
6262
exc_type: PyObjectRef,
6363
exc_val: OptionalArg,
6464
exc_tb: OptionalArg,
6565
vm: &VirtualMachine,
66-
) -> PyResult {
67-
self.inner.throw(
66+
) -> PyResult<PyIterReturn> {
67+
zelf.inner.throw(
68+
zelf.as_object(),
6869
exc_type,
6970
exc_val.unwrap_or_none(vm),
7071
exc_tb.unwrap_or_none(vm),
@@ -73,8 +74,8 @@ impl PyGenerator {
7374
}
7475

7576
#[pymethod]
76-
fn close(&self, vm: &VirtualMachine) -> PyResult<()> {
77-
self.inner.close(vm)
77+
fn close(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<()> {
78+
zelf.inner.close(zelf.as_object(), vm)
7879
}
7980

8081
#[pyproperty]
@@ -98,8 +99,7 @@ impl PyGenerator {
9899
impl IteratorIterable for PyGenerator {}
99100
impl SlotIterator for PyGenerator {
100101
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
101-
// TODO: Fix zelf.send to return PyIterReturn
102-
PyIterReturn::from_result(zelf.send(vm.ctx.none(), vm), vm)
102+
Self::send(zelf.clone(), vm.ctx.none(), vm)
103103
}
104104
}
105105

‎vm/src/builtins/map.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ impl SlotIterator for PyMap {
5858
}
5959

6060
// the mapper itself can raise StopIteration which does stop the map iteration
61-
PyIterReturn::from_result(vm.invoke(&zelf.mapper, next_objs), vm)
61+
PyIterReturn::from_pyresult(vm.invoke(&zelf.mapper, next_objs), vm)
6262
}
6363
}
6464

‎vm/src/builtins/pytype.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -263,7 +263,7 @@ impl PyType {
263263
}
264264
"__next__" => {
265265
let func: slots::IterNextFunc = |zelf, vm| {
266-
PyIterReturn::from_result(
266+
PyIterReturn::from_pyresult(
267267
vm.call_special_method(zelf.clone(), "__next__", ()),
268268
vm,
269269
)

0 commit comments

Comments
 (0)