From 577cea92fe1a2eeac44d64bc681ff2ed1b6a1d4e Mon Sep 17 00:00:00 2001 From: coolreader18 <33094578+coolreader18@users.noreply.github.com> Date: Wed, 16 Oct 2019 20:07:30 -0500 Subject: [PATCH] Add coroutines, async/await functionality, and gen.close() --- bytecode/src/bytecode.rs | 3 + compiler/src/compile.rs | 40 ++++++++--- vm/src/builtins.rs | 1 + vm/src/exceptions.rs | 3 + vm/src/frame.rs | 38 +++++++--- vm/src/obj/mod.rs | 1 + vm/src/obj/objcoroutine.rs | 137 +++++++++++++++++++++++++++++++++++++ vm/src/obj/objgenerator.rs | 47 ++++++++++--- vm/src/types.rs | 8 +++ vm/src/vm.rs | 3 + 10 files changed, 255 insertions(+), 26 deletions(-) create mode 100644 vm/src/obj/objcoroutine.rs diff --git a/bytecode/src/bytecode.rs b/bytecode/src/bytecode.rs index f6ea86d2fb..8475b88feb 100644 --- a/bytecode/src/bytecode.rs +++ b/bytecode/src/bytecode.rs @@ -55,6 +55,7 @@ bitflags! { const HAS_ANNOTATIONS = 0x04; const NEW_LOCALS = 0x08; const IS_GENERATOR = 0x10; + const IS_COROUTINE = 0x20; } } @@ -273,6 +274,7 @@ pub enum Instruction { Reverse { amount: usize, }, + GetAwaitable, } use self::Instruction::*; @@ -549,6 +551,7 @@ impl Instruction { FormatValue { spec, .. } => w!(FormatValue, spec), // TODO: write conversion PopException => w!(PopException), Reverse { amount } => w!(Reverse, amount), + GetAwaitable => w!(GetAwaitable), } } } diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs index ac4f00a659..f467cc5c3c 100644 --- a/compiler/src/compile.rs +++ b/compiler/src/compile.rs @@ -29,6 +29,7 @@ struct Compiler { current_qualified_path: Option, in_loop: bool, in_function_def: bool, + in_async_func: bool, optimize: u8, } @@ -126,6 +127,7 @@ impl Compiler { current_qualified_path: None, in_loop: false, in_function_def: false, + in_async_func: false, optimize, } } @@ -451,11 +453,7 @@ impl Compiler { decorator_list, returns, } => { - if *is_async { - unimplemented!("async def"); - } else { - self.compile_function_def(name, args, body, decorator_list, returns)? - } + self.compile_function_def(name, args, body, decorator_list, returns, *is_async)?; } ClassDef { name, @@ -797,11 +795,15 @@ impl Compiler { body: &[ast::Statement], decorator_list: &[ast::Expression], returns: &Option, // TODO: use type hint somehow.. + is_async: bool, ) -> Result<(), CompileError> { // Create bytecode for this function: // remember to restore self.in_loop to the original after the function is compiled let was_in_loop = self.in_loop; let was_in_function_def = self.in_function_def; + + let was_in_async_func = self.in_async_func; + self.in_async_func = is_async; self.in_loop = false; self.in_function_def = true; @@ -870,6 +872,10 @@ impl Compiler { }); } + if is_async { + code.flags |= bytecode::CodeFlags::IS_COROUTINE; + } + self.emit(Instruction::LoadConst { value: bytecode::Constant::Code { code: Box::new(code), @@ -891,6 +897,7 @@ impl Compiler { self.current_qualified_path = old_qualified_path; self.in_loop = was_in_loop; self.in_function_def = was_in_function_def; + self.in_async_func = was_in_async_func; Ok(()) } @@ -1551,7 +1558,7 @@ impl Compiler { self.emit(Instruction::BuildSlice { size }); } Yield { value } => { - if !self.in_function_def { + if !self.in_function_def || self.in_async_func { return Err(CompileError { error: CompileErrorType::InvalidYield, location: self.current_source_location.clone(), @@ -1566,8 +1573,13 @@ impl Compiler { }; self.emit(Instruction::YieldValue); } - Await { .. } => { - unimplemented!("await"); + Await { value } => { + self.compile_expression(value)?; + self.emit(Instruction::GetAwaitable); + self.emit(Instruction::LoadConst { + value: bytecode::Constant::None, + }); + self.emit(Instruction::YieldFrom); } YieldFrom { value } => { self.mark_generator(); @@ -1612,8 +1624,14 @@ impl Compiler { self.load_name(name); } Lambda { args, body } => { + let was_in_loop = self.in_loop; + let was_in_function_def = self.in_function_def; + let was_in_async_func = self.in_async_func; + self.in_async_func = false; + self.in_loop = false; + self.in_function_def = true; + let name = "".to_string(); - // no need to worry about the self.loop_depth because there are no loops in lambda expressions self.enter_function(&name, args)?; self.compile_expression(body)?; self.emit(Instruction::ReturnValue); @@ -1629,6 +1647,10 @@ impl Compiler { }); // Turn code object into function object: self.emit(Instruction::MakeFunction); + + self.in_loop = was_in_loop; + self.in_function_def = was_in_function_def; + self.in_async_func = was_in_async_func; } Comprehension { kind, generators } => { self.compile_comprehension(kind, generators)?; diff --git a/vm/src/builtins.rs b/vm/src/builtins.rs index 82af90fd99..6f0d68626e 100644 --- a/vm/src/builtins.rs +++ b/vm/src/builtins.rs @@ -900,6 +900,7 @@ pub fn make_module(vm: &VirtualMachine, module: PyObjectRef) { "UserWarning" => ctx.exceptions.user_warning.clone(), "KeyboardInterrupt" => ctx.exceptions.keyboard_interrupt.clone(), + "GeneratorExit" => ctx.exceptions.generator_exit.clone(), "SystemExit" => ctx.exceptions.system_exit.clone(), }); } diff --git a/vm/src/exceptions.rs b/vm/src/exceptions.rs index 43682a8d9d..da6c479dab 100644 --- a/vm/src/exceptions.rs +++ b/vm/src/exceptions.rs @@ -271,6 +271,7 @@ pub struct ExceptionZoo { pub user_warning: PyClassRef, pub keyboard_interrupt: PyClassRef, + pub generator_exit: PyClassRef, pub system_exit: PyClassRef, } @@ -327,6 +328,7 @@ impl ExceptionZoo { let user_warning = create_type("UserWarning", &type_type, &warning); let keyboard_interrupt = create_type("KeyboardInterrupt", &type_type, &base_exception_type); + let generator_exit = create_type("GeneratorExit", &type_type, &base_exception_type); let system_exit = create_type("SystemExit", &type_type, &base_exception_type); ExceptionZoo { @@ -376,6 +378,7 @@ impl ExceptionZoo { reference_error, user_warning, keyboard_interrupt, + generator_exit, system_exit, } } diff --git a/vm/src/frame.rs b/vm/src/frame.rs index 923a901827..65690e32e2 100644 --- a/vm/src/frame.rs +++ b/vm/src/frame.rs @@ -8,16 +8,16 @@ use crate::bytecode; use crate::function::PyFuncArgs; use crate::obj::objbool; use crate::obj::objcode::PyCodeRef; +use crate::obj::objcoroutine::PyCoroutine; use crate::obj::objdict::{PyDict, PyDictRef}; +use crate::obj::objgenerator::PyGenerator; use crate::obj::objiter; use crate::obj::objlist; use crate::obj::objslice::PySlice; -use crate::obj::objstr; -use crate::obj::objstr::PyString; +use crate::obj::objstr::{self, PyString}; use crate::obj::objtraceback::{PyTraceback, PyTracebackRef}; use crate::obj::objtuple::PyTuple; -use crate::obj::objtype; -use crate::obj::objtype::PyClassRef; +use crate::obj::objtype::{self, PyClassRef}; use crate::pyobject::{ IdProtocol, ItemProtocol, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol, }; @@ -230,7 +230,7 @@ impl Frame { exc_type: PyClassRef, exc_val: PyObjectRef, exc_tb: PyObjectRef, - ) -> PyResult { + ) -> PyResult { if let bytecode::Instruction::YieldFrom = self.code.instructions[self.lasti.get()] { let coro = self.last_value(); vm.call_method( @@ -243,7 +243,8 @@ impl Frame { self.lasti.set(self.lasti.get() + 1); let val = objiter::stop_iter_value(vm, &err)?; self._send(coro, val, vm) - }) + })?; + Ok(ExecutionResult::Return(vm.get_none())) } else { let exception = vm.new_exception_obj(exc_type, vec![exc_val])?; match self.unwind_blocks(vm, UnwindReason::Raising { exception }) { @@ -251,7 +252,6 @@ impl Frame { Ok(Some(result)) => Ok(result), Err(exception) => Err(exception), } - .and_then(|res| res.into_result(vm)) } } @@ -465,6 +465,24 @@ impl Frame { self.push_value(iter_obj); Ok(None) } + bytecode::Instruction::GetAwaitable => { + let awaited_obj = self.pop_value(); + let awaitable = if awaited_obj.payload_is::() + { + awaited_obj + } else { + let await_method = + vm.get_method_or_type_error(awaited_obj.clone(), "__await__", || { + format!( + "object {} can't be used in 'await' expression", + awaited_obj.class().name, + ) + })?; + vm.invoke(&await_method, vec![])? + }; + self.push_value(awaitable); + Ok(None) + } bytecode::Instruction::ForIter { target } => self.execute_for_iter(vm, *target), bytecode::Instruction::MakeFunction => self.execute_make_function(vm), bytecode::Instruction::CallFunction { typ } => self.execute_call_function(vm, typ), @@ -1016,7 +1034,11 @@ impl Frame { } fn _send(&self, coro: PyObjectRef, val: PyObjectRef, vm: &VirtualMachine) -> PyResult { - if vm.is_none(&val) { + 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]) diff --git a/vm/src/obj/mod.rs b/vm/src/obj/mod.rs index d527d69d4f..57233025da 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 objcoroutine; pub mod objdict; pub mod objellipsis; pub mod objenumerate; diff --git a/vm/src/obj/objcoroutine.rs b/vm/src/obj/objcoroutine.rs new file mode 100644 index 0000000000..e052b401b7 --- /dev/null +++ b/vm/src/obj/objcoroutine.rs @@ -0,0 +1,137 @@ +use super::objtype::{isinstance, issubclass, PyClassRef}; +use crate::frame::{ExecutionResult, FrameRef}; +use crate::function::OptionalArg; +use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; +use crate::vm::VirtualMachine; + +pub type PyCoroutineRef = PyRef; + +#[pyclass(name = "coroutine")] +#[derive(Debug)] +pub struct PyCoroutine { + frame: FrameRef, +} + +impl PyValue for PyCoroutine { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.ctx.types.coroutine_type.clone() + } +} + +#[pyimpl] +impl PyCoroutine { + pub fn new(frame: FrameRef, vm: &VirtualMachine) -> PyCoroutineRef { + PyCoroutine { frame }.into_ref(vm) + } + + #[pymethod] + pub(crate) fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult { + self.frame.push_value(value.clone()); + + vm.run_frame(self.frame.clone())?.into_result(vm) + } + + #[pymethod] + fn throw( + &self, + exc_type: PyClassRef, + exc_val: OptionalArg, + exc_tb: OptionalArg, + 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 !issubclass(&exc_type, &vm.ctx.exceptions.base_exception_type) { + return Err(vm.new_type_error("Can't throw non exception".to_string())); + } + vm.frames.borrow_mut().push(self.frame.clone()); + let result = self + .frame + .gen_throw( + vm, + exc_type, + exc_val.unwrap_or(vm.get_none()), + exc_tb.unwrap_or(vm.get_none()), + ) + .and_then(|res| res.into_result(vm)); + vm.frames.borrow_mut().pop(); + result + } + + #[pymethod] + fn close(&self, vm: &VirtualMachine) -> PyResult<()> { + vm.frames.borrow_mut().push(self.frame.clone()); + let result = self.frame.gen_throw( + vm, + vm.ctx.exceptions.generator_exit.clone(), + vm.get_none(), + vm.get_none(), + ); + vm.frames.borrow_mut().pop(); + match result { + Ok(ExecutionResult::Yield(_)) => Err(vm.new_exception( + vm.ctx.exceptions.runtime_error.clone(), + "generator ignored GeneratorExit".to_string(), + )), + Err(e) => { + if isinstance(&e, &vm.ctx.exceptions.generator_exit) { + Ok(()) + } else { + Err(e) + } + } + _ => Ok(()), + } + } + + #[pymethod(name = "__await__")] + fn r#await(zelf: PyRef, _vm: &VirtualMachine) -> PyCoroutineWrapper { + PyCoroutineWrapper { coro: zelf } + } +} + +#[pyclass(name = "coroutine_wrapper")] +#[derive(Debug)] +pub struct PyCoroutineWrapper { + coro: PyCoroutineRef, +} + +impl PyValue for PyCoroutineWrapper { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.ctx.types.coroutine_wrapper_type.clone() + } +} + +#[pyimpl] +impl PyCoroutineWrapper { + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef, _vm: &VirtualMachine) -> PyRef { + zelf + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + self.coro.send(vm.get_none(), vm) + } + + #[pymethod] + fn send(&self, val: PyObjectRef, vm: &VirtualMachine) -> PyResult { + self.coro.send(val, vm) + } + + #[pymethod] + fn throw( + &self, + exc_type: PyClassRef, + exc_val: OptionalArg, + exc_tb: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult { + self.coro.throw(exc_type, exc_val, exc_tb, vm) + } +} + +pub fn init(ctx: &PyContext) { + PyCoroutine::extend_class(ctx, &ctx.types.coroutine_type); + PyCoroutineWrapper::extend_class(ctx, &ctx.types.coroutine_wrapper_type); +} diff --git a/vm/src/obj/objgenerator.rs b/vm/src/obj/objgenerator.rs index 11159e0556..44342bfed3 100644 --- a/vm/src/obj/objgenerator.rs +++ b/vm/src/obj/objgenerator.rs @@ -2,8 +2,8 @@ * The mythical generator. */ -use super::objtype::{issubclass, PyClassRef}; -use crate::frame::FrameRef; +use super::objtype::{isinstance, issubclass, PyClassRef}; +use crate::frame::{ExecutionResult, FrameRef}; use crate::function::OptionalArg; use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; use crate::vm::VirtualMachine; @@ -39,7 +39,7 @@ impl PyGenerator { } #[pymethod] - fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult { + pub(crate) fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult { self.frame.push_value(value.clone()); vm.run_frame(self.frame.clone())?.into_result(vm) @@ -59,15 +59,44 @@ impl PyGenerator { return Err(vm.new_type_error("Can't throw non exception".to_string())); } vm.frames.borrow_mut().push(self.frame.clone()); - let result = self.frame.gen_throw( - vm, - exc_type, - exc_val.unwrap_or(vm.get_none()), - exc_tb.unwrap_or(vm.get_none()), - ); + let result = self + .frame + .gen_throw( + vm, + exc_type, + exc_val.unwrap_or(vm.get_none()), + exc_tb.unwrap_or(vm.get_none()), + ) + .and_then(|res| res.into_result(vm)); vm.frames.borrow_mut().pop(); result } + + #[pymethod] + fn close(&self, vm: &VirtualMachine) -> PyResult<()> { + vm.frames.borrow_mut().push(self.frame.clone()); + let result = self.frame.gen_throw( + vm, + vm.ctx.exceptions.generator_exit.clone(), + vm.get_none(), + vm.get_none(), + ); + vm.frames.borrow_mut().pop(); + match result { + Ok(ExecutionResult::Yield(_)) => Err(vm.new_exception( + vm.ctx.exceptions.runtime_error.clone(), + "generator ignored GeneratorExit".to_string(), + )), + Err(e) => { + if isinstance(&e, &vm.ctx.exceptions.generator_exit) { + Ok(()) + } else { + Err(e) + } + } + _ => Ok(()), + } + } } pub fn init(ctx: &PyContext) { diff --git a/vm/src/types.rs b/vm/src/types.rs index 52612af464..8c3bddcd43 100644 --- a/vm/src/types.rs +++ b/vm/src/types.rs @@ -4,6 +4,7 @@ use crate::obj::objbytes; use crate::obj::objclassmethod; use crate::obj::objcode; use crate::obj::objcomplex; +use crate::obj::objcoroutine; use crate::obj::objdict; use crate::obj::objellipsis; use crate::obj::objenumerate; @@ -50,6 +51,8 @@ pub struct TypeZoo { pub bool_type: PyClassRef, pub classmethod_type: PyClassRef, pub code_type: PyClassRef, + pub coroutine_type: PyClassRef, + pub coroutine_wrapper_type: PyClassRef, pub dict_type: PyClassRef, pub enumerate_type: PyClassRef, pub filter_type: PyClassRef, @@ -122,6 +125,8 @@ impl TypeZoo { let weakref_type = create_type("ref", &type_type, &object_type); let weakproxy_type = create_type("weakproxy", &type_type, &object_type); let generator_type = create_type("generator", &type_type, &object_type); + let coroutine_type = create_type("coroutine", &type_type, &object_type); + let coroutine_wrapper_type = create_type("coroutine_wrapper", &type_type, &object_type); let bound_method_type = create_type("method", &type_type, &object_type); let str_type = create_type("str", &type_type, &object_type); let list_type = create_type("list", &type_type, &object_type); @@ -170,6 +175,8 @@ impl TypeZoo { bytes_type, bytesiterator_type, code_type, + coroutine_type, + coroutine_wrapper_type, complex_type, classmethod_type, int_type, @@ -291,6 +298,7 @@ pub fn initialize_types(context: &PyContext) { objstaticmethod::init(&context); objclassmethod::init(&context); objgenerator::init(&context); + objcoroutine::init(&context); objint::init(&context); objfloat::init(&context); objcomplex::init(&context); diff --git a/vm/src/vm.rs b/vm/src/vm.rs index d7d07df531..403d975921 100644 --- a/vm/src/vm.rs +++ b/vm/src/vm.rs @@ -26,6 +26,7 @@ use crate::import; use crate::obj::objbool; use crate::obj::objbuiltinfunc::PyBuiltinFunction; use crate::obj::objcode::{PyCode, PyCodeRef}; +use crate::obj::objcoroutine::PyCoroutine; use crate::obj::objdict::PyDictRef; use crate::obj::objfunction::{PyFunction, PyMethod}; use crate::obj::objgenerator::PyGenerator; @@ -664,6 +665,8 @@ impl VirtualMachine { // If we have a generator, create a new generator if code.flags.contains(bytecode::CodeFlags::IS_GENERATOR) { Ok(PyGenerator::new(frame, self).into_object()) + } else if code.flags.contains(bytecode::CodeFlags::IS_COROUTINE) { + Ok(PyCoroutine::new(frame, self).into_object()) } else { self.run_frame_full(frame) }