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 0d4ee3462..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,15 +1021,19 @@ impl Frame { Err(exception) } + 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, + }) + } + fn _send(&self, coro: PyObjectRef, val: PyObjectRef, vm: &VirtualMachine) -> PyResult { - if let Some(gen) = coro.payload::() { - gen.send(val, vm) - } else if let Some(coro) = coro.payload::() { - coro.send(val, vm) - } else if vm.is_none(&val) { - objiter::call_next(vm, &coro) - } else { - vm.call_method(&coro, "send", vec![val]) + 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/macros.rs b/vm/src/macros.rs index 4b2ad0b04..5eeb824a3 100644 --- a/vm/src/macros.rs +++ b/vm/src/macros.rs @@ -244,6 +244,10 @@ macro_rules! match_class { let $binding = $obj; $default }}; + (match ($obj:expr) { ref $binding:ident => $default:expr $(,)? }) => {{ + let $binding = &$obj; + $default + }}; // An arm taken when the object is an instance of the specified built-in // class and binding the downcasted object to the specified identifier and @@ -251,6 +255,9 @@ macro_rules! match_class { (match ($obj:expr) { $binding:ident @ $class:ty => $expr:block $($rest:tt)* }) => { $crate::match_class!(match ($obj) { $binding @ $class => ($expr), $($rest)* }) }; + (match ($obj:expr) { ref $binding:ident @ $class:ty => $expr:block $($rest:tt)* }) => { + $crate::match_class!(match ($obj) { ref $binding @ $class => ($expr), $($rest)* }) + }; // An arm taken when the object is an instance of the specified built-in // class and binding the downcasted object to the specified identifier. @@ -260,6 +267,12 @@ macro_rules! match_class { Err(_obj) => $crate::match_class!(match (_obj) { $($rest)* }), } }; + (match ($obj:expr) { ref $binding:ident @ $class:ty => $expr:expr, $($rest:tt)* }) => { + match $obj.payload::<$class>() { + Some($binding) => $expr, + None => $crate::match_class!(match ($obj) { $($rest)* }), + } + }; // An arm taken when the object is an instance of the specified built-in // class and the target expression is a block. diff --git a/vm/src/obj/mod.rs b/vm/src/obj/mod.rs index af5d531d1..dd259c752 100644 --- a/vm/src/obj/mod.rs +++ b/vm/src/obj/mod.rs @@ -8,6 +8,7 @@ pub mod objbytes; pub mod objclassmethod; pub mod objcode; pub mod objcomplex; +pub mod objcoroinner; pub mod objcoroutine; pub mod objdict; pub mod objellipsis; diff --git a/vm/src/obj/objcoroinner.rs b/vm/src/obj/objcoroinner.rs new file mode 100644 index 000000000..e544e96fe --- /dev/null +++ b/vm/src/obj/objcoroinner.rs @@ -0,0 +1,101 @@ +use super::{objiter, objtype}; +use crate::exceptions::{self, PyBaseExceptionRef}; +use crate::frame::{ExecutionResult, FrameRef}; +use crate::pyobject::{PyObjectRef, PyResult}; +use crate::vm::VirtualMachine; + +use std::cell::Cell; + +#[derive(Debug)] +pub struct Coro { + frame: FrameRef, + closed: Cell, + running: Cell, +} + +impl Coro { + pub fn new(frame: FrameRef) -> Self { + Coro { + frame, + closed: Cell::new(false), + running: Cell::new(false), + } + } + + fn maybe_close(&self, res: &PyResult) { + match res { + Ok(ExecutionResult::Return(_)) | Err(_) => self.closed.set(true), + Ok(ExecutionResult::Yield(_)) => {} + } + } + + pub fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult { + if self.closed.get() { + return Err(objiter::new_stop_iteration(vm)); + } + + 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) + } + + pub fn throw( + &self, + exc_type: PyObjectRef, + exc_val: PyObjectRef, + exc_tb: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + if self.closed.get() { + 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) + } + + pub fn close(&self, vm: &VirtualMachine) -> PyResult<()> { + if self.closed.get() { + 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 { + Ok(ExecutionResult::Yield(_)) => { + Err(vm.new_runtime_error("generator ignored GeneratorExit".to_owned())) + } + Err(e) if !is_gen_exit(&e, vm) => Err(e), + _ => 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 { + objtype::isinstance(exc, &vm.ctx.exceptions.generator_exit) +} diff --git a/vm/src/obj/objcoroutine.rs b/vm/src/obj/objcoroutine.rs index 6f30d8eb6..ac7a77fb2 100644 --- a/vm/src/obj/objcoroutine.rs +++ b/vm/src/obj/objcoroutine.rs @@ -1,20 +1,18 @@ -use super::objiter::new_stop_iteration; -use super::objtype::{isinstance, PyClassRef}; -use crate::exceptions; -use crate::frame::{ExecutionResult, FrameRef}; +use super::objcode::PyCodeRef; +use super::objcoroinner::Coro; +use super::objstr::PyStringRef; +use super::objtype::PyClassRef; +use crate::frame::FrameRef; use crate::function::OptionalArg; use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; use crate::vm::VirtualMachine; -use std::cell::Cell; - pub type PyCoroutineRef = PyRef; #[pyclass(name = "coroutine")] #[derive(Debug)] pub struct PyCoroutine { - frame: FrameRef, - closed: Cell, + inner: Coro, } impl PyValue for PyCoroutine { @@ -25,33 +23,20 @@ impl PyValue for PyCoroutine { #[pyimpl] impl PyCoroutine { + pub fn as_coro(&self) -> &Coro { + &self.inner + } + pub fn new(frame: FrameRef, vm: &VirtualMachine) -> PyCoroutineRef { PyCoroutine { - frame, - closed: Cell::new(false), + inner: Coro::new(frame), } .into_ref(vm) } - // TODO: deduplicate this code with objgenerator - fn maybe_close(&self, res: &PyResult) { - match res { - Ok(ExecutionResult::Return(_)) | Err(_) => self.closed.set(true), - Ok(ExecutionResult::Yield(_)) => {} - } - } - #[pymethod] - pub(crate) fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult { - if self.closed.get() { - return Err(new_stop_iteration(vm)); - } - - self.frame.push_value(value.clone()); - - let result = vm.run_frame(self.frame.clone()); - self.maybe_close(&result); - result?.into_result(vm) + fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult { + self.inner.send(value, vm) } #[pymethod] @@ -62,51 +47,46 @@ impl PyCoroutine { exc_tb: OptionalArg, vm: &VirtualMachine, ) -> PyResult { - let exc_val = exc_val.unwrap_or_else(|| vm.get_none()); - let exc_tb = exc_tb.unwrap_or_else(|| vm.get_none()); - if self.closed.get() { - return Err(exceptions::normalize(exc_type, exc_val, exc_tb, vm)?); - } - vm.frames.borrow_mut().push(self.frame.clone()); - let result = self.frame.gen_throw(vm, exc_type, exc_val, exc_tb); - self.maybe_close(&result); - vm.frames.borrow_mut().pop(); - result?.into_result(vm) + self.inner.throw( + exc_type, + exc_val.unwrap_or_else(|| vm.get_none()), + exc_tb.unwrap_or_else(|| vm.get_none()), + vm, + ) } #[pymethod] fn close(&self, vm: &VirtualMachine) -> PyResult<()> { - if self.closed.get() { - return Ok(()); - } - vm.frames.borrow_mut().push(self.frame.clone()); - let result = self.frame.gen_throw( - vm, - vm.ctx.exceptions.generator_exit.clone().into_object(), - vm.get_none(), - vm.get_none(), - ); - vm.frames.borrow_mut().pop(); - self.closed.set(true); - match result { - Ok(ExecutionResult::Yield(_)) => { - Err(vm.new_runtime_error("generator ignored GeneratorExit".to_owned())) - } - Err(e) => { - if isinstance(&e, &vm.ctx.exceptions.generator_exit) { - Ok(()) - } else { - Err(e) - } - } - _ => Ok(()), - } + self.inner.close(vm) } #[pymethod(name = "__await__")] 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 4e4256a7f..45b81de99 100644 --- a/vm/src/obj/objgenerator.rs +++ b/vm/src/obj/objgenerator.rs @@ -2,23 +2,20 @@ * The mythical generator. */ -use super::objiter::new_stop_iteration; -use super::objtype::{isinstance, PyClassRef}; -use crate::exceptions; -use crate::frame::{ExecutionResult, FrameRef}; +use super::objcode::PyCodeRef; +use super::objcoroinner::Coro; +use super::objtype::PyClassRef; +use crate::frame::FrameRef; use crate::function::OptionalArg; use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; use crate::vm::VirtualMachine; -use std::cell::Cell; - pub type PyGeneratorRef = PyRef; #[pyclass(name = "generator")] #[derive(Debug)] pub struct PyGenerator { - frame: FrameRef, - closed: Cell, + inner: Coro, } impl PyValue for PyGenerator { @@ -29,19 +26,15 @@ impl PyValue for PyGenerator { #[pyimpl] impl PyGenerator { - pub fn new(frame: FrameRef, vm: &VirtualMachine) -> PyGeneratorRef { - PyGenerator { - frame, - closed: Cell::new(false), - } - .into_ref(vm) + pub fn as_coro(&self) -> &Coro { + &self.inner } - fn maybe_close(&self, res: &PyResult) { - match res { - Ok(ExecutionResult::Return(_)) | Err(_) => self.closed.set(true), - Ok(ExecutionResult::Yield(_)) => {} + pub fn new(frame: FrameRef, vm: &VirtualMachine) -> PyGeneratorRef { + PyGenerator { + inner: Coro::new(frame), } + .into_ref(vm) } #[pymethod(name = "__iter__")] @@ -55,16 +48,8 @@ impl PyGenerator { } #[pymethod] - pub(crate) fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult { - if self.closed.get() { - return Err(new_stop_iteration(vm)); - } - - self.frame.push_value(value.clone()); - - let result = vm.run_frame(self.frame.clone()); - self.maybe_close(&result); - result?.into_result(vm) + fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult { + self.inner.send(value, vm) } #[pymethod] @@ -75,45 +60,34 @@ impl PyGenerator { exc_tb: OptionalArg, vm: &VirtualMachine, ) -> PyResult { - let exc_val = exc_val.unwrap_or_else(|| vm.get_none()); - let exc_tb = exc_tb.unwrap_or_else(|| vm.get_none()); - if self.closed.get() { - return Err(exceptions::normalize(exc_type, exc_val, exc_tb, vm)?); - } - vm.frames.borrow_mut().push(self.frame.clone()); - let result = self.frame.gen_throw(vm, exc_type, exc_val, exc_tb); - self.maybe_close(&result); - vm.frames.borrow_mut().pop(); - result?.into_result(vm) + self.inner.throw( + exc_type, + exc_val.unwrap_or_else(|| vm.get_none()), + exc_tb.unwrap_or_else(|| vm.get_none()), + vm, + ) } #[pymethod] fn close(&self, vm: &VirtualMachine) -> PyResult<()> { - if self.closed.get() { - return Ok(()); - } - vm.frames.borrow_mut().push(self.frame.clone()); - let result = self.frame.gen_throw( - vm, - vm.ctx.exceptions.generator_exit.clone().into_object(), - vm.get_none(), - vm.get_none(), - ); - vm.frames.borrow_mut().pop(); - self.closed.set(true); - match result { - Ok(ExecutionResult::Yield(_)) => { - Err(vm.new_runtime_error("generator ignored GeneratorExit".to_owned())) - } - Err(e) => { - if isinstance(&e, &vm.ctx.exceptions.generator_exit) { - Ok(()) - } else { - Err(e) - } - } - _ => Ok(()), - } + 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() } } diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs index 8ea94e842..786e1a6bf 100644 --- a/vm/src/pyobject.rs +++ b/vm/src/pyobject.rs @@ -640,11 +640,9 @@ impl PyObject { /// another downcast can be attempted without unnecessary cloning. pub fn downcast(self: Rc) -> Result, PyObjectRef> { if self.payload_is::() { - Ok({ - PyRef { - obj: self, - _payload: PhantomData, - } + Ok(PyRef { + obj: self, + _payload: PhantomData, }) } else { Err(self)