Merge pull request RustPython#820 from skinny121/gen_throw

Implement generator.throw

Author: windelbouwman

tests/snippets/generators.py

@@ -1,3 +1,5 @@
+from testutils import assertRaises
+
 
 r = []
 
@@ -42,3 +44,23 @@ def g4():
 r = list(g4())
 assert r == [None, (2,)]
 
+
+def catch_exception():
+    try:
+        yield 1
+    except ValueError:
+        yield 2
+        yield 3
+
+
+g = catch_exception()
+assert next(g) == 1
+
+assert g.throw(ValueError, ValueError(), None) == 2
+assert next(g) == 3
+
+g = catch_exception()
+assert next(g) == 1
+
+with assertRaises(KeyError):
+    assert g.throw(KeyError, KeyError(), None) == 2

vm/src/frame.rs

@@ -283,6 +283,17 @@ impl Frame {
         }
     }
 
+    pub fn throw(
+        &self,
+        vm: &VirtualMachine,
+        exception: PyObjectRef,
+    ) -> Result<ExecutionResult, PyObjectRef> {
+        match self.unwind_exception(vm, exception) {
+            None => self.run(vm),
+            Some(exception) => Err(exception),
+        }
+    }
+
     pub fn fetch_instruction(&self) -> &bytecode::Instruction {
         let ins2 = &self.code.instructions[*self.lasti.borrow()];
         *self.lasti.borrow_mut() += 1;

vm/src/obj/objgenerator.rs

@@ -3,12 +3,13 @@
  */
 
 use crate::frame::{ExecutionResult, FrameRef};
-use crate::obj::objtype::PyClassRef;
-use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue};
+use crate::obj::objtype::{isinstance, PyClassRef};
+use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
 use crate::vm::VirtualMachine;
 
 pub type PyGeneratorRef = PyRef<PyGenerator>;
 
+#[pyclass(name = "generator", __inside_vm)]
 #[derive(Debug)]
 pub struct PyGenerator {
     frame: FrameRef,
@@ -20,38 +21,59 @@ impl PyValue for PyGenerator {
     }
 }
 
-impl PyGeneratorRef {
+#[pyimpl(__inside_vm)]
+impl PyGenerator {
     pub fn new(frame: FrameRef, vm: &VirtualMachine) -> PyGeneratorRef {
         PyGenerator { frame }.into_ref(vm)
     }
 
-    fn iter(self, _vm: &VirtualMachine) -> PyGeneratorRef {
-        self
+    #[pymethod(name = "__iter__")]
+    fn iter(zelf: PyGeneratorRef, _vm: &VirtualMachine) -> PyGeneratorRef {
+        zelf
     }
 
-    fn next(self, vm: &VirtualMachine) -> PyResult {
+    #[pymethod(name = "__next__")]
+    fn next(&self, vm: &VirtualMachine) -> PyResult {
         self.send(vm.get_none(), vm)
     }
 
-    fn send(self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+    #[pymethod]
+    fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult {
         self.frame.push_value(value.clone());
 
-        match vm.run_frame(self.frame.clone())? {
-            ExecutionResult::Yield(value) => Ok(value),
-            ExecutionResult::Return(_value) => {
-                // Stop iteration!
-                let stop_iteration = vm.ctx.exceptions.stop_iteration.clone();
-                Err(vm.new_exception(stop_iteration, "End of generator".to_string()))
-            }
+        let result = vm.run_frame(self.frame.clone())?;
+        handle_execution_result(result, vm)
+    }
+
+    #[pymethod]
+    fn throw(
+        &self,
+        _exc_type: PyObjectRef,
+        exc_val: PyObjectRef,
+        _exc_tb: PyObjectRef,
+        vm: &VirtualMachine,
+    ) -> PyResult {
+        // TODO what should we do with the other parameters? CPython normalises them with
+        //      PyErr_NormalizeException, do we want to do the same.
+        if !isinstance(&exc_val, &vm.ctx.exceptions.base_exception_type) {
+            return Err(vm.new_type_error("Can't throw non exception".to_string()));
+        }
+        let result = vm.frame_throw(self.frame.clone(), exc_val)?;
+        handle_execution_result(result, vm)
+    }
+}
+
+fn handle_execution_result(result: ExecutionResult, vm: &VirtualMachine) -> PyResult {
+    match result {
+        ExecutionResult::Yield(value) => Ok(value),
+        ExecutionResult::Return(_value) => {
+            // Stop iteration!
+            let stop_iteration = vm.ctx.exceptions.stop_iteration.clone();
+            Err(vm.new_exception(stop_iteration, "End of generator".to_string()))
         }
     }
 }
 
-pub fn init(context: &PyContext) {
-    let generator_type = &context.generator_type;
-    extend_class!(context, generator_type, {
-        "__iter__" => context.new_rustfunc(PyGeneratorRef::iter),
-        "__next__" => context.new_rustfunc(PyGeneratorRef::next),
-        "send" => context.new_rustfunc(PyGeneratorRef::send)
-    });
+pub fn init(ctx: &PyContext) {
+    PyGenerator::extend_class(ctx, &ctx.generator_type);
 }

vm/src/vm.rs

@@ -21,7 +21,7 @@ use crate::obj::objbuiltinfunc::PyBuiltinFunction;
 use crate::obj::objcode::PyCodeRef;
 use crate::obj::objdict::PyDictRef;
 use crate::obj::objfunction::{PyFunction, PyMethod};
-use crate::obj::objgenerator::PyGeneratorRef;
+use crate::obj::objgenerator::PyGenerator;
 use crate::obj::objiter;
 use crate::obj::objsequence;
 use crate::obj::objstr::{PyString, PyStringRef};
@@ -94,6 +94,17 @@ impl VirtualMachine {
         result
     }
 
+    pub fn frame_throw(
+        &self,
+        frame: FrameRef,
+        exception: PyObjectRef,
+    ) -> PyResult<ExecutionResult> {
+        self.frames.borrow_mut().push(frame.clone());
+        let result = frame.throw(self, exception);
+        self.frames.borrow_mut().pop();
+        result
+    }
+
     pub fn current_frame(&self) -> Option<Ref<FrameRef>> {
         let frames = self.frames.borrow();
         if frames.is_empty() {
@@ -374,7 +385,7 @@ impl VirtualMachine {
 
         // If we have a generator, create a new generator
         if code.code.is_generator {
-            Ok(PyGeneratorRef::new(frame, self).into_object())
+            Ok(PyGenerator::new(frame, self).into_object())
         } else {
             self.run_frame_full(frame)
         }