From 43713c41e4b0d7617bfbe46a83c7ef844c551c19 Mon Sep 17 00:00:00 2001 From: coolreader18 <33094578+coolreader18@users.noreply.github.com> Date: Sat, 4 Jan 2020 15:06:36 -0600 Subject: [PATCH] Add gi_ and cr_ attributes to generators and coroutines --- tests/snippets/generators.py | 27 ++++++++++++++++++++++ vm/src/frame.rs | 43 ++++++++++++++++++++++++------------ vm/src/obj/objcoroinner.rs | 19 +++++++++++++++- vm/src/obj/objcoroutine.rs | 25 +++++++++++++++++++++ vm/src/obj/objgenerator.rs | 18 +++++++++++++++ 5 files changed, 117 insertions(+), 15 deletions(-) diff --git a/tests/snippets/generators.py b/tests/snippets/generators.py index dccbff3ec..b2d994755 100644 --- a/tests/snippets/generators.py +++ b/tests/snippets/generators.py @@ -141,3 +141,30 @@ next(g) assert_raises(TypeError, g.throw, TypeError) assert_raises(StopIteration, next, g) assert_raises(TypeError, g.throw, TypeError) + +def a(): + assert g.gi_running + try: + yield + except: + assert g.gi_running + + +g = a() +next(g) +assert_raises(StopIteration, g.throw, TypeError) + +g = a() +next(g) +g.close() + +it = iter([1,2,3,4]) + +def a(): + yield from it + +g = a() +assert next(g) == 1 +assert g.gi_yieldfrom is it +assert list(g) == [2,3,4] +assert g.gi_yieldfrom is None diff --git a/vm/src/frame.rs b/vm/src/frame.rs index 5ce2db57e..681d5b2d4 100644 --- a/vm/src/frame.rs +++ b/vm/src/frame.rs @@ -9,6 +9,7 @@ use crate::exceptions::{self, ExceptionCtor, PyBaseExceptionRef}; use crate::function::{single_or_tuple_any, PyFuncArgs}; use crate::obj::objbool; use crate::obj::objcode::PyCodeRef; +use crate::obj::objcoroinner::Coro; use crate::obj::objcoroutine::PyCoroutine; use crate::obj::objdict::{PyDict, PyDictRef}; use crate::obj::objgenerator::PyGenerator; @@ -205,6 +206,15 @@ impl Frame { } } + pub fn yield_from_target(&self) -> Option { + if let Some(bytecode::Instruction::YieldFrom) = self.code.instructions.get(self.lasti.get()) + { + Some(self.last_value()) + } else { + None + } + } + pub(crate) fn gen_throw( &self, vm: &VirtualMachine, @@ -212,16 +222,18 @@ impl Frame { exc_val: PyObjectRef, exc_tb: PyObjectRef, ) -> PyResult { - if let bytecode::Instruction::YieldFrom = self.code.instructions[self.lasti.get()] { - let coro = self.last_value(); - vm.call_method(&coro, "throw", vec![exc_type, exc_val, exc_tb]) - .or_else(|err| { - self.pop_value(); - self.lasti.set(self.lasti.get() + 1); - let val = objiter::stop_iter_value(vm, &err)?; - self._send(coro, val, vm) - }) - .map(ExecutionResult::Yield) + if let Some(coro) = self.yield_from_target() { + let res = match self.builtin_coro(&coro) { + Some(coro) => coro.throw(exc_type, exc_val, exc_tb, vm), + None => vm.call_method(&coro, "throw", vec![exc_type, exc_val, exc_tb]), + }; + res.or_else(|err| { + self.pop_value(); + self.lasti.set(self.lasti.get() + 1); + let val = objiter::stop_iter_value(vm, &err)?; + self._send(coro, val, vm) + }) + .map(ExecutionResult::Yield) } else { let exception = exceptions::normalize(exc_type, exc_val, exc_tb, vm)?; match self.unwind_blocks(vm, UnwindReason::Raising { exception }) { @@ -1009,13 +1021,16 @@ impl Frame { Err(exception) } - fn _send(&self, coro: PyObjectRef, val: PyObjectRef, vm: &VirtualMachine) -> PyResult { - let builtin_coro = match_class!(match coro { + fn builtin_coro<'a>(&self, coro: &'a PyObjectRef) -> Option<&'a Coro> { + match_class!(match coro { ref g @ PyGenerator => Some(g.as_coro()), ref c @ PyCoroutine => Some(c.as_coro()), _ => None, - }); - match builtin_coro { + }) + } + + fn _send(&self, coro: PyObjectRef, val: PyObjectRef, vm: &VirtualMachine) -> PyResult { + match self.builtin_coro(&coro) { Some(coro) => coro.send(val, vm), None if vm.is_none(&val) => objiter::call_next(vm, &coro), None => vm.call_method(&coro, "send", vec![val]), diff --git a/vm/src/obj/objcoroinner.rs b/vm/src/obj/objcoroinner.rs index 529b30159..e544e96fe 100644 --- a/vm/src/obj/objcoroinner.rs +++ b/vm/src/obj/objcoroinner.rs @@ -10,6 +10,7 @@ use std::cell::Cell; pub struct Coro { frame: FrameRef, closed: Cell, + running: Cell, } impl Coro { @@ -17,6 +18,7 @@ impl Coro { Coro { frame, closed: Cell::new(false), + running: Cell::new(false), } } @@ -33,8 +35,9 @@ impl Coro { } self.frame.push_value(value.clone()); - + self.running.set(true); let result = vm.run_frame(self.frame.clone()); + self.running.set(false); self.maybe_close(&result); result?.into_result(vm) } @@ -50,7 +53,9 @@ impl Coro { return Err(exceptions::normalize(exc_type, exc_val, exc_tb, vm)?); } vm.frames.borrow_mut().push(self.frame.clone()); + self.running.set(true); let result = self.frame.gen_throw(vm, exc_type, exc_val, exc_tb); + self.running.set(false); self.maybe_close(&result); vm.frames.borrow_mut().pop(); result?.into_result(vm) @@ -61,12 +66,14 @@ impl Coro { return Ok(()); } vm.frames.borrow_mut().push(self.frame.clone()); + self.running.set(true); let result = self.frame.gen_throw( vm, vm.ctx.exceptions.generator_exit.clone().into_object(), vm.get_none(), vm.get_none(), ); + self.running.set(false); vm.frames.borrow_mut().pop(); self.closed.set(true); match result { @@ -77,6 +84,16 @@ impl Coro { _ => Ok(()), } } + + pub fn closed(&self) -> bool { + self.closed.get() + } + pub fn running(&self) -> bool { + self.running.get() + } + pub fn frame(&self) -> FrameRef { + self.frame.clone() + } } pub fn is_gen_exit(exc: &PyBaseExceptionRef, vm: &VirtualMachine) -> bool { diff --git a/vm/src/obj/objcoroutine.rs b/vm/src/obj/objcoroutine.rs index d82cb8859..ac7a77fb2 100644 --- a/vm/src/obj/objcoroutine.rs +++ b/vm/src/obj/objcoroutine.rs @@ -1,4 +1,6 @@ +use super::objcode::PyCodeRef; use super::objcoroinner::Coro; +use super::objstr::PyStringRef; use super::objtype::PyClassRef; use crate::frame::FrameRef; use crate::function::OptionalArg; @@ -62,6 +64,29 @@ impl PyCoroutine { fn r#await(zelf: PyRef) -> PyCoroutineWrapper { PyCoroutineWrapper { coro: zelf } } + + #[pyproperty] + fn cr_await(&self, _vm: &VirtualMachine) -> Option { + self.inner.frame().yield_from_target() + } + #[pyproperty] + fn cr_frame(&self, _vm: &VirtualMachine) -> FrameRef { + self.inner.frame() + } + #[pyproperty] + fn cr_running(&self, _vm: &VirtualMachine) -> bool { + self.inner.running() + } + #[pyproperty] + fn cr_code(&self, _vm: &VirtualMachine) -> PyCodeRef { + self.inner.frame().code.clone() + } + // TODO: coroutine origin tracking: + // https://docs.python.org/3/library/sys.html#sys.set_coroutine_origin_tracking_depth + #[pyproperty] + fn cr_origin(&self, _vm: &VirtualMachine) -> Option<(PyStringRef, usize, PyStringRef)> { + None + } } #[pyclass(name = "coroutine_wrapper")] diff --git a/vm/src/obj/objgenerator.rs b/vm/src/obj/objgenerator.rs index 19c526e48..45b81de99 100644 --- a/vm/src/obj/objgenerator.rs +++ b/vm/src/obj/objgenerator.rs @@ -2,6 +2,7 @@ * The mythical generator. */ +use super::objcode::PyCodeRef; use super::objcoroinner::Coro; use super::objtype::PyClassRef; use crate::frame::FrameRef; @@ -71,6 +72,23 @@ impl PyGenerator { fn close(&self, vm: &VirtualMachine) -> PyResult<()> { self.inner.close(vm) } + + #[pyproperty] + fn gi_frame(&self, _vm: &VirtualMachine) -> FrameRef { + self.inner.frame() + } + #[pyproperty] + fn gi_running(&self, _vm: &VirtualMachine) -> bool { + self.inner.running() + } + #[pyproperty] + fn gi_code(&self, _vm: &VirtualMachine) -> PyCodeRef { + self.inner.frame().code.clone() + } + #[pyproperty] + fn gi_yieldfrom(&self, _vm: &VirtualMachine) -> Option { + self.inner.frame().yield_from_target() + } } pub fn init(ctx: &PyContext) {