diff --git a/bytecode/src/bytecode.rs b/bytecode/src/bytecode.rs index 8475b88fe..4e760bba2 100644 --- a/bytecode/src/bytecode.rs +++ b/bytecode/src/bytecode.rs @@ -218,9 +218,8 @@ pub enum Instruction { SetupWith { end: Label, }, - CleanupWith { - end: Label, - }, + WithCleanupStart, + WithCleanupFinish, PopBlock, Raise { argc: usize, @@ -275,6 +274,10 @@ pub enum Instruction { amount: usize, }, GetAwaitable, + BeforeAsyncWith, + SetupAsyncWith { + end: Label, + }, } use self::Instruction::*; @@ -528,7 +531,10 @@ impl Instruction { EnterFinally => w!(EnterFinally), EndFinally => w!(EndFinally), SetupWith { end } => w!(SetupWith, label_map[end]), - CleanupWith { end } => w!(CleanupWith, label_map[end]), + WithCleanupStart => w!(WithCleanupStart), + WithCleanupFinish => w!(WithCleanupFinish), + BeforeAsyncWith => w!(BeforeAsyncWith), + SetupAsyncWith { end } => w!(SetupAsyncWith, label_map[end]), PopBlock => w!(PopBlock), Raise { argc } => w!(Raise, argc), BuildString { size } => w!(BuildString, size), diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs index f467cc5c3..fae828dcd 100644 --- a/compiler/src/compile.rs +++ b/compiler/src/compile.rs @@ -386,27 +386,72 @@ impl Compiler { body, } => { if *is_async { - unimplemented!("async with"); - } else { - let end_label = self.new_label(); - for item in items { - self.compile_expression(&item.context_expr)?; - self.emit(Instruction::SetupWith { end: end_label }); - match &item.optional_vars { - Some(var) => { - self.compile_store(var)?; + let end_labels = items + .iter() + .map(|item| { + let end_label = self.new_label(); + self.compile_expression(&item.context_expr)?; + self.emit(Instruction::BeforeAsyncWith); + self.emit(Instruction::GetAwaitable); + self.emit(Instruction::LoadConst { + value: bytecode::Constant::None, + }); + self.emit(Instruction::YieldFrom); + self.emit(Instruction::SetupAsyncWith { end: end_label }); + match &item.optional_vars { + Some(var) => { + self.compile_store(var)?; + } + None => { + self.emit(Instruction::Pop); + } } - None => { - self.emit(Instruction::Pop); - } - } - } + Ok(end_label) + }) + .collect::, CompileError>>()?; self.compile_statements(body)?; - for _ in 0..items.len() { - self.emit(Instruction::CleanupWith { end: end_label }); + + for end_label in end_labels { + self.emit(Instruction::PopBlock); + self.emit(Instruction::EnterFinally); + self.set_label(end_label); + self.emit(Instruction::WithCleanupStart); + self.emit(Instruction::GetAwaitable); + self.emit(Instruction::LoadConst { + value: bytecode::Constant::None, + }); + self.emit(Instruction::YieldFrom); + self.emit(Instruction::WithCleanupFinish); + } + } else { + let end_labels = items + .iter() + .map(|item| { + let end_label = self.new_label(); + self.compile_expression(&item.context_expr)?; + self.emit(Instruction::SetupWith { end: end_label }); + match &item.optional_vars { + Some(var) => { + self.compile_store(var)?; + } + None => { + self.emit(Instruction::Pop); + } + } + Ok(end_label) + }) + .collect::, CompileError>>()?; + + self.compile_statements(body)?; + + for end_label in end_labels { + self.emit(Instruction::PopBlock); + self.emit(Instruction::EnterFinally); + self.set_label(end_label); + self.emit(Instruction::WithCleanupStart); + self.emit(Instruction::WithCleanupFinish); } - self.set_label(end_label); } } For { diff --git a/vm/src/builtins.rs b/vm/src/builtins.rs index 6f0d68626..5cf2db1af 100644 --- a/vm/src/builtins.rs +++ b/vm/src/builtins.rs @@ -873,6 +873,7 @@ pub fn make_module(vm: &VirtualMachine, module: PyObjectRef) { "FileNotFoundError" => ctx.exceptions.file_not_found_error.clone(), "FileExistsError" => ctx.exceptions.file_exists_error.clone(), "StopIteration" => ctx.exceptions.stop_iteration.clone(), + "StopAsyncIteration" => ctx.exceptions.stop_async_iteration.clone(), "SystemError" => ctx.exceptions.system_error.clone(), "PermissionError" => ctx.exceptions.permission_error.clone(), "UnicodeError" => ctx.exceptions.unicode_error.clone(), diff --git a/vm/src/exceptions.rs b/vm/src/exceptions.rs index da6c479da..cc2f5cad3 100644 --- a/vm/src/exceptions.rs +++ b/vm/src/exceptions.rs @@ -244,6 +244,7 @@ pub struct ExceptionZoo { pub reference_error: PyClassRef, pub runtime_error: PyClassRef, pub stop_iteration: PyClassRef, + pub stop_async_iteration: PyClassRef, pub syntax_error: PyClassRef, pub indentation_error: PyClassRef, pub tab_error: PyClassRef, @@ -292,6 +293,7 @@ impl ExceptionZoo { let runtime_error = create_type("RuntimeError", &type_type, &exception_type); let reference_error = create_type("ReferenceError", &type_type, &exception_type); let stop_iteration = create_type("StopIteration", &type_type, &exception_type); + let stop_async_iteration = create_type("StopAsyncIteration", &type_type, &exception_type); let syntax_error = create_type("SyntaxError", &type_type, &exception_type); let system_error = create_type("SystemError", &type_type, &exception_type); let type_error = create_type("TypeError", &type_type, &exception_type); @@ -352,6 +354,7 @@ impl ExceptionZoo { permission_error, runtime_error, stop_iteration, + stop_async_iteration, syntax_error, indentation_error, tab_error, diff --git a/vm/src/frame.rs b/vm/src/frame.rs index 65690e32e..dfdf5697f 100644 --- a/vm/src/frame.rs +++ b/vm/src/frame.rs @@ -49,10 +49,6 @@ enum BlockType { FinallyHandler { reason: Option, }, - With { - end: bytecode::Label, - context_manager: PyObjectRef, - }, ExceptHandler, } @@ -431,30 +427,70 @@ impl Frame { } bytecode::Instruction::SetupWith { end } => { let context_manager = self.pop_value(); + let exit = vm.get_attribute(context_manager.clone(), "__exit__")?; + self.push_value(exit); // Call enter: - let obj = vm.call_method(&context_manager, "__enter__", vec![])?; - self.push_block(BlockType::With { - end: *end, - context_manager: context_manager.clone(), - }); - self.push_value(obj); + let enter_res = vm.call_method(&context_manager, "__enter__", vec![])?; + self.push_block(BlockType::Finally { handler: *end }); + self.push_value(enter_res); Ok(None) } - bytecode::Instruction::CleanupWith { end: end1 } => { - let block = self.pop_block(); - if let BlockType::With { - end: end2, - context_manager, - } = &block.typ - { - debug_assert!(end1 == end2); - self.call_context_manager_exit(vm, &context_manager, None)?; - } else { - unreachable!("Block stack is incorrect, expected a with block"); - } + bytecode::Instruction::BeforeAsyncWith => { + let mgr = self.pop_value(); + let aexit = vm.get_attribute(mgr.clone(), "__aexit__")?; + self.push_value(aexit); + let aenter_res = vm.call_method(&mgr, "__aenter__", vec![])?; + self.push_value(aenter_res); Ok(None) } + bytecode::Instruction::SetupAsyncWith { end } => { + self.push_block(BlockType::Finally { handler: *end }); + Ok(None) + } + bytecode::Instruction::WithCleanupStart => { + let block = self.current_block().unwrap(); + let reason = match block.typ { + BlockType::FinallyHandler { reason } => reason, + _ => panic!("WithCleanupStart expects a FinallyHandler block on stack"), + }; + let exc = reason.and_then(|reason| match reason { + UnwindReason::Raising { exception } => Some(exception), + _ => None, + }); + + let exit = self.pop_value(); + + let args = if let Some(exc) = exc { + let exc_type = exc.class().into_object(); + let exc_val = exc.clone(); + let exc_tb = vm.ctx.none(); // TODO: retrieve traceback? + vec![exc_type, exc_val, exc_tb] + } else { + vec![vm.ctx.none(), vm.ctx.none(), vm.ctx.none()] + }; + let exit_res = vm.invoke(&exit, args)?; + self.push_value(exit_res); + + Ok(None) + } + bytecode::Instruction::WithCleanupFinish => { + let block = self.pop_block(); + let reason = match block.typ { + BlockType::FinallyHandler { reason } => reason, + _ => panic!("WithCleanupFinish expects a FinallyHandler block on stack"), + }; + + let suppress_exception = objbool::boolval(vm, self.pop_value())?; + if suppress_exception { + // suppress exception + Ok(None) + } else if let Some(reason) = reason { + self.unwind_blocks(vm, reason) + } else { + Ok(None) + } + } bytecode::Instruction::PopBlock => { self.pop_block(); Ok(None) @@ -663,7 +699,7 @@ impl Frame { /// unwinding a block. /// Optionally returns an exception. #[cfg_attr(feature = "flame-it", flame("Frame"))] - fn unwind_blocks(&self, vm: &VirtualMachine, mut reason: UnwindReason) -> FrameResult { + fn unwind_blocks(&self, vm: &VirtualMachine, reason: UnwindReason) -> FrameResult { // First unwind all existing blocks on the block stack: while let Some(block) = self.current_block() { match block.typ { @@ -699,60 +735,6 @@ impl Frame { return Ok(None); } } - BlockType::With { - context_manager, - end, - } => { - self.pop_block(); - match &reason { - UnwindReason::Raising { exception } => { - match self.call_context_manager_exit( - vm, - &context_manager, - Some(exception.clone()), - ) { - Ok(exit_result_obj) => { - match objbool::boolval(vm, exit_result_obj) { - // If __exit__ method returned True, suppress the exception and continue execution. - Ok(suppress_exception) => { - if suppress_exception { - self.jump(end); - return Ok(None); - } else { - // go on with the stack unwinding. - } - } - Err(exit_exc) => { - reason = UnwindReason::Raising { - exception: exit_exc, - }; - // return Err(exit_exc); - } - } - } - Err(exit_exc) => { - // TODO: what about original exception? - reason = UnwindReason::Raising { - exception: exit_exc, - }; - // return Err(exit_exc); - } - } - } - _ => { - match self.call_context_manager_exit(vm, &context_manager, None) { - Ok(..) => {} - Err(exit_exc) => { - // __exit__ went wrong, - reason = UnwindReason::Raising { - exception: exit_exc, - }; - // return Err(exc); - } - } - } - } - } BlockType::FinallyHandler { .. } => { self.pop_block(); } @@ -773,26 +755,6 @@ impl Frame { } } - fn call_context_manager_exit( - &self, - vm: &VirtualMachine, - context_manager: &PyObjectRef, - exc: Option, - ) -> PyResult { - // TODO: do we want to put the exit call on the stack? - // TODO: what happens when we got an error during execution of __exit__? - let (exc_type, exc_val, exc_tb) = if let Some(exc) = exc { - let exc_type = exc.class().into_object(); - let exc_val = exc.clone(); - let exc_tb = vm.ctx.none(); // TODO: retrieve traceback? - (exc_type, exc_val, exc_tb) - } else { - (vm.ctx.none(), vm.ctx.none(), vm.ctx.none()) - }; - - vm.call_method(context_manager, "__exit__", vec![exc_type, exc_val, exc_tb]) - } - fn store_name( &self, vm: &VirtualMachine,