Add async with, reorganize how with blocks work

Author: coolreader18

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),

compiler/src/compile.rs

@@ -386,27 +386,72 @@ impl<O: OutputStream> Compiler<O> {
                 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)?;
-                            }
-                            None => {
-                                self.emit(Instruction::Pop);
+                    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);
+                                }
                             }
-                        }
+                            Ok(end_label)
+                        })
+                        .collect::<Result<Vec<_>, 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::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::<Result<Vec<_>, 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::WithCleanupFinish);
                     }
-                    self.set_label(end_label);
                 }
             }
             For {

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(),

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,

vm/src/frame.rs

@@ -49,10 +49,6 @@ enum BlockType {
     FinallyHandler {
         reason: Option<UnwindReason>,
     },
-    With {
-        end: bytecode::Label,
-        context_manager: PyObjectRef,
-    },
     ExceptHandler,
 }
 
@@ -431,29 +427,69 @@ 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(),
+                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::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,
                 });
-                self.push_value(obj);
+
+                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::CleanupWith { end: end1 } => {
+            bytecode::Instruction::WithCleanupFinish => {
                 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)?;
+                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 {
-                    unreachable!("Block stack is incorrect, expected a with block");
+                    Ok(None)
                 }
-
-                Ok(None)
             }
             bytecode::Instruction::PopBlock => {
                 self.pop_block();
@@ -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<PyObjectRef>,
-    ) -> 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,