Add gi_ and cr_ attributes to generators and coroutines

Author: coolreader18

tests/snippets/generators.py

@@ -141,3 +141,30 @@ def a():
 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

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,23 +206,34 @@ impl Frame {
         }
     }
 
+    pub fn yield_from_target(&self) -> Option<PyObjectRef> {
+        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,
         exc_type: PyObjectRef,
         exc_val: PyObjectRef,
         exc_tb: PyObjectRef,
     ) -> PyResult<ExecutionResult> {
-        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]),

vm/src/obj/objcoroinner.rs

@@ -10,13 +10,15 @@ use std::cell::Cell;
 pub struct Coro {
     frame: FrameRef,
     closed: Cell<bool>,
+    running: Cell<bool>,
 }
 
 impl Coro {
     pub fn new(frame: FrameRef) -> Self {
         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 {

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<Self>) -> PyCoroutineWrapper {
         PyCoroutineWrapper { coro: zelf }
     }
+
+    #[pyproperty]
+    fn cr_await(&self, _vm: &VirtualMachine) -> Option<PyObjectRef> {
+        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")]

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<PyObjectRef> {
+        self.inner.frame().yield_from_target()
+    }
 }
 
 pub fn init(ctx: &PyContext) {