Skip to content

Commit 43713c4

Browse files
coolreader18coolreader18
committed
Add gi_ and cr_ attributes to generators and coroutines
1 parent 48de33c commit 43713c4

5 files changed

Lines changed: 117 additions & 15 deletions

File tree

‎tests/snippets/generators.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,3 +141,30 @@ def a():
141141
assert_raises(TypeError, g.throw, TypeError)
142142
assert_raises(StopIteration, next, g)
143143
assert_raises(TypeError, g.throw, TypeError)
144+
145+
def a():
146+
assert g.gi_running
147+
try:
148+
yield
149+
except:
150+
assert g.gi_running
151+
152+
153+
g = a()
154+
next(g)
155+
assert_raises(StopIteration, g.throw, TypeError)
156+
157+
g = a()
158+
next(g)
159+
g.close()
160+
161+
it = iter([1,2,3,4])
162+
163+
def a():
164+
yield from it
165+
166+
g = a()
167+
assert next(g) == 1
168+
assert g.gi_yieldfrom is it
169+
assert list(g) == [2,3,4]
170+
assert g.gi_yieldfrom is None

‎vm/src/frame.rs‎

Lines changed: 29 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use crate::exceptions::{self, ExceptionCtor, PyBaseExceptionRef};
99
use crate::function::{single_or_tuple_any, PyFuncArgs};
1010
use crate::obj::objbool;
1111
use crate::obj::objcode::PyCodeRef;
12+
use crate::obj::objcoroinner::Coro;
1213
use crate::obj::objcoroutine::PyCoroutine;
1314
use crate::obj::objdict::{PyDict, PyDictRef};
1415
use crate::obj::objgenerator::PyGenerator;
@@ -205,23 +206,34 @@ impl Frame {
205206
}
206207
}
207208

209+
pub fn yield_from_target(&self) -> Option<PyObjectRef> {
210+
if let Some(bytecode::Instruction::YieldFrom) = self.code.instructions.get(self.lasti.get())
211+
{
212+
Some(self.last_value())
213+
} else {
214+
None
215+
}
216+
}
217+
208218
pub(crate) fn gen_throw(
209219
&self,
210220
vm: &VirtualMachine,
211221
exc_type: PyObjectRef,
212222
exc_val: PyObjectRef,
213223
exc_tb: PyObjectRef,
214224
) -> PyResult<ExecutionResult> {
215-
if let bytecode::Instruction::YieldFrom = self.code.instructions[self.lasti.get()] {
216-
let coro = self.last_value();
217-
vm.call_method(&coro, "throw", vec![exc_type, exc_val, exc_tb])
218-
.or_else(|err| {
219-
self.pop_value();
220-
self.lasti.set(self.lasti.get() + 1);
221-
let val = objiter::stop_iter_value(vm, &err)?;
222-
self._send(coro, val, vm)
223-
})
224-
.map(ExecutionResult::Yield)
225+
if let Some(coro) = self.yield_from_target() {
226+
let res = match self.builtin_coro(&coro) {
227+
Some(coro) => coro.throw(exc_type, exc_val, exc_tb, vm),
228+
None => vm.call_method(&coro, "throw", vec![exc_type, exc_val, exc_tb]),
229+
};
230+
res.or_else(|err| {
231+
self.pop_value();
232+
self.lasti.set(self.lasti.get() + 1);
233+
let val = objiter::stop_iter_value(vm, &err)?;
234+
self._send(coro, val, vm)
235+
})
236+
.map(ExecutionResult::Yield)
225237
} else {
226238
let exception = exceptions::normalize(exc_type, exc_val, exc_tb, vm)?;
227239
match self.unwind_blocks(vm, UnwindReason::Raising { exception }) {
@@ -1009,13 +1021,16 @@ impl Frame {
10091021
Err(exception)
10101022
}
10111023

1012-
fn _send(&self, coro: PyObjectRef, val: PyObjectRef, vm: &VirtualMachine) -> PyResult {
1013-
let builtin_coro = match_class!(match coro {
1024+
fn builtin_coro<'a>(&self, coro: &'a PyObjectRef) -> Option<&'a Coro> {
1025+
match_class!(match coro {
10141026
ref g @ PyGenerator => Some(g.as_coro()),
10151027
ref c @ PyCoroutine => Some(c.as_coro()),
10161028
_ => None,
1017-
});
1018-
match builtin_coro {
1029+
})
1030+
}
1031+
1032+
fn _send(&self, coro: PyObjectRef, val: PyObjectRef, vm: &VirtualMachine) -> PyResult {
1033+
match self.builtin_coro(&coro) {
10191034
Some(coro) => coro.send(val, vm),
10201035
None if vm.is_none(&val) => objiter::call_next(vm, &coro),
10211036
None => vm.call_method(&coro, "send", vec![val]),

‎vm/src/obj/objcoroinner.rs‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,15 @@ use std::cell::Cell;
1010
pub struct Coro {
1111
frame: FrameRef,
1212
closed: Cell<bool>,
13+
running: Cell<bool>,
1314
}
1415

1516
impl Coro {
1617
pub fn new(frame: FrameRef) -> Self {
1718
Coro {
1819
frame,
1920
closed: Cell::new(false),
21+
running: Cell::new(false),
2022
}
2123
}
2224

@@ -33,8 +35,9 @@ impl Coro {
3335
}
3436

3537
self.frame.push_value(value.clone());
36-
38+
self.running.set(true);
3739
let result = vm.run_frame(self.frame.clone());
40+
self.running.set(false);
3841
self.maybe_close(&result);
3942
result?.into_result(vm)
4043
}
@@ -50,7 +53,9 @@ impl Coro {
5053
return Err(exceptions::normalize(exc_type, exc_val, exc_tb, vm)?);
5154
}
5255
vm.frames.borrow_mut().push(self.frame.clone());
56+
self.running.set(true);
5357
let result = self.frame.gen_throw(vm, exc_type, exc_val, exc_tb);
58+
self.running.set(false);
5459
self.maybe_close(&result);
5560
vm.frames.borrow_mut().pop();
5661
result?.into_result(vm)
@@ -61,12 +66,14 @@ impl Coro {
6166
return Ok(());
6267
}
6368
vm.frames.borrow_mut().push(self.frame.clone());
69+
self.running.set(true);
6470
let result = self.frame.gen_throw(
6571
vm,
6672
vm.ctx.exceptions.generator_exit.clone().into_object(),
6773
vm.get_none(),
6874
vm.get_none(),
6975
);
76+
self.running.set(false);
7077
vm.frames.borrow_mut().pop();
7178
self.closed.set(true);
7279
match result {
@@ -77,6 +84,16 @@ impl Coro {
7784
_ => Ok(()),
7885
}
7986
}
87+
88+
pub fn closed(&self) -> bool {
89+
self.closed.get()
90+
}
91+
pub fn running(&self) -> bool {
92+
self.running.get()
93+
}
94+
pub fn frame(&self) -> FrameRef {
95+
self.frame.clone()
96+
}
8097
}
8198

8299
pub fn is_gen_exit(exc: &PyBaseExceptionRef, vm: &VirtualMachine) -> bool {

‎vm/src/obj/objcoroutine.rs‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1+
use super::objcode::PyCodeRef;
12
use super::objcoroinner::Coro;
3+
use super::objstr::PyStringRef;
24
use super::objtype::PyClassRef;
35
use crate::frame::FrameRef;
46
use crate::function::OptionalArg;
@@ -62,6 +64,29 @@ impl PyCoroutine {
6264
fn r#await(zelf: PyRef<Self>) -> PyCoroutineWrapper {
6365
PyCoroutineWrapper { coro: zelf }
6466
}
67+
68+
#[pyproperty]
69+
fn cr_await(&self, _vm: &VirtualMachine) -> Option<PyObjectRef> {
70+
self.inner.frame().yield_from_target()
71+
}
72+
#[pyproperty]
73+
fn cr_frame(&self, _vm: &VirtualMachine) -> FrameRef {
74+
self.inner.frame()
75+
}
76+
#[pyproperty]
77+
fn cr_running(&self, _vm: &VirtualMachine) -> bool {
78+
self.inner.running()
79+
}
80+
#[pyproperty]
81+
fn cr_code(&self, _vm: &VirtualMachine) -> PyCodeRef {
82+
self.inner.frame().code.clone()
83+
}
84+
// TODO: coroutine origin tracking:
85+
// https://docs.python.org/3/library/sys.html#sys.set_coroutine_origin_tracking_depth
86+
#[pyproperty]
87+
fn cr_origin(&self, _vm: &VirtualMachine) -> Option<(PyStringRef, usize, PyStringRef)> {
88+
None
89+
}
6590
}
6691

6792
#[pyclass(name = "coroutine_wrapper")]

‎vm/src/obj/objgenerator.rs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
* The mythical generator.
33
*/
44

5+
use super::objcode::PyCodeRef;
56
use super::objcoroinner::Coro;
67
use super::objtype::PyClassRef;
78
use crate::frame::FrameRef;
@@ -71,6 +72,23 @@ impl PyGenerator {
7172
fn close(&self, vm: &VirtualMachine) -> PyResult<()> {
7273
self.inner.close(vm)
7374
}
75+
76+
#[pyproperty]
77+
fn gi_frame(&self, _vm: &VirtualMachine) -> FrameRef {
78+
self.inner.frame()
79+
}
80+
#[pyproperty]
81+
fn gi_running(&self, _vm: &VirtualMachine) -> bool {
82+
self.inner.running()
83+
}
84+
#[pyproperty]
85+
fn gi_code(&self, _vm: &VirtualMachine) -> PyCodeRef {
86+
self.inner.frame().code.clone()
87+
}
88+
#[pyproperty]
89+
fn gi_yieldfrom(&self, _vm: &VirtualMachine) -> Option<PyObjectRef> {
90+
self.inner.frame().yield_from_target()
91+
}
7492
}
7593

7694
pub fn init(ctx: &PyContext) {

0 commit comments

Comments
 (0)