RuntimeError on coro raising StopIteration

Author: coolreader18

Lib/test/test_asyncgen.py

@@ -219,7 +219,6 @@ async def gen():
                                     'async generator.*StopAsyncIteration'):
             to_list(gen())
 
-    @unittest.skip("TODO: RUSTPYTHON")
     def test_async_gen_exception_06(self):
         async def gen():
             yield 123

vm/src/obj/objasyncgenerator.rs

@@ -1,5 +1,5 @@
 use super::objcode::PyCodeRef;
-use super::objcoroinner::Coro;
+use super::objcoroinner::{Coro, Variant};
 use super::objtype::{self, PyClassRef};
 use crate::frame::FrameRef;
 use crate::function::OptionalArg;
@@ -30,7 +30,7 @@ impl PyAsyncGen {
 
     pub fn new(frame: FrameRef, vm: &VirtualMachine) -> PyAsyncGenRef {
         PyAsyncGen {
-            inner: Coro::new_async(frame),
+            inner: Coro::new(frame, Variant::AsyncGen),
             running_async: Cell::new(false),
         }
         .into_ref(vm)

vm/src/obj/objcoroinner.rs

@@ -6,35 +6,44 @@ use crate::vm::VirtualMachine;
 
 use std::cell::{Cell, RefCell};
 
+#[derive(Debug, PartialEq, Clone, Copy)]
+pub enum Variant {
+    Gen,
+    Coroutine,
+    AsyncGen,
+}
+impl Variant {
+    fn exec_result(self, res: ExecutionResult, vm: &VirtualMachine) -> PyResult {
+        res.into_result(self == Self::AsyncGen, vm)
+    }
+    fn name(self) -> &'static str {
+        match self {
+            Self::Gen => "generator",
+            Self::Coroutine => "coroutine",
+            Self::AsyncGen => "async generator",
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct Coro {
     frame: FrameRef,
     pub closed: Cell<bool>,
     running: Cell<bool>,
     exceptions: RefCell<Vec<PyBaseExceptionRef>>,
     started: Cell<bool>,
-    async_iter: bool,
+    variant: Variant,
 }
 
 impl Coro {
-    pub fn new(frame: FrameRef) -> Self {
+    pub fn new(frame: FrameRef, variant: Variant) -> Self {
         Coro {
             frame,
             closed: Cell::new(false),
             running: Cell::new(false),
             exceptions: RefCell::new(vec![]),
             started: Cell::new(false),
-            async_iter: false,
-        }
-    }
-    pub fn new_async(frame: FrameRef) -> Self {
-        Coro {
-            frame,
-            closed: Cell::new(false),
-            running: Cell::new(false),
-            exceptions: RefCell::new(vec![]),
-            started: Cell::new(false),
-            async_iter: true,
+            variant,
         }
     }
 
@@ -62,6 +71,22 @@ impl Coro {
         );
         self.running.set(false);
         self.started.set(true);
+        if let Err(ref e) = result {
+            if objtype::isinstance(e, &vm.ctx.exceptions.stop_iteration) {
+                let err =
+                    vm.new_runtime_error(format!("{} raised StopIteration", self.variant.name()));
+                err.set_cause(Some(e.clone()));
+                return Err(err);
+            }
+            if self.variant == Variant::AsyncGen
+                && objtype::isinstance(e, &vm.ctx.exceptions.stop_async_iteration)
+            {
+                let err =
+                    vm.new_runtime_error("async generator raise StopAsyncIteration".to_owned());
+                err.set_cause(Some(e.clone()));
+                return Err(err);
+            }
+        }
         result
     }
 
@@ -77,7 +102,7 @@ impl Coro {
         self.frame.push_value(value.clone());
         let result = self.run_with_context(|| vm.run_frame(self.frame.clone()), vm);
         self.maybe_close(&result);
-        result?.into_result(self.async_iter, vm)
+        self.variant.exec_result(result?, vm)
     }
     pub fn throw(
         &self,
@@ -94,7 +119,7 @@ impl Coro {
             self.run_with_context(|| self.frame.gen_throw(vm, exc_type, exc_val, exc_tb), vm);
         self.maybe_close(&result);
         vm.frames.borrow_mut().pop();
-        result?.into_result(self.async_iter, vm)
+        self.variant.exec_result(result?, vm)
     }
 
     pub fn close(&self, vm: &VirtualMachine) -> PyResult<()> {

vm/src/obj/objcoroutine.rs

@@ -1,5 +1,5 @@
 use super::objcode::PyCodeRef;
-use super::objcoroinner::Coro;
+use super::objcoroinner::{Coro, Variant};
 use super::objstr::PyStringRef;
 use super::objtype::PyClassRef;
 use crate::frame::FrameRef;
@@ -29,7 +29,7 @@ impl PyCoroutine {
 
     pub fn new(frame: FrameRef, vm: &VirtualMachine) -> PyCoroutineRef {
         PyCoroutine {
-            inner: Coro::new(frame),
+            inner: Coro::new(frame, Variant::Coroutine),
         }
         .into_ref(vm)
     }

vm/src/obj/objgenerator.rs

@@ -3,7 +3,7 @@
  */
 
 use super::objcode::PyCodeRef;
-use super::objcoroinner::Coro;
+use super::objcoroinner::{Coro, Variant};
 use super::objtype::PyClassRef;
 use crate::frame::FrameRef;
 use crate::function::OptionalArg;
@@ -32,7 +32,7 @@ impl PyGenerator {
 
     pub fn new(frame: FrameRef, vm: &VirtualMachine) -> PyGeneratorRef {
         PyGenerator {
-            inner: Coro::new(frame),
+            inner: Coro::new(frame, Variant::Gen),
         }
         .into_ref(vm)
     }