Fix with __exit__ error handling

youknowone · youknowone · commit 929d9a82bc6c · 2024-04-25T14:48:47.000+09:00
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -925,15 +925,15 @@ impl ExecutingFrame<'_> {
                     _ => None,
                 };
 
-                let exit = self.pop_value();
+                let exit = self.top_value();
 
                 let args = if let Some(exc) = exc {
                     vm.split_exception(exc)
                 } else {
                     (vm.ctx.none(), vm.ctx.none(), vm.ctx.none())
                 };
                 let exit_res = exit.call(args, vm)?;
-                self.push_value(exit_res);
+                self.replace_top(exit_res);
 
                 Ok(None)
             }
@@ -1954,6 +1954,13 @@ impl ExecutingFrame<'_> {
         self.state.stack.drain(stack_len - count..)
     }
 
+    #[inline]
+    fn replace_top(&mut self, mut top: PyObjectRef) -> PyObjectRef {
+        let last = self.state.stack.last_mut().unwrap();
+        std::mem::swap(&mut top, last);
+        top
+    }
+
     #[inline]
     #[track_caller] // not a real track_caller but top_value is not very useful
     fn top_value(&self) -> &PyObject {
