From 93c2608c197c7e4843b19ca2e4dce33a12e146e1 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@foxmail.com>
Date: Fri, 6 Dec 2024 13:22:02 +0800
Subject: [PATCH] fix context bug

---
 src/compiler/compiler.c |  1 +
 src/interpreter/ceval.c |  1 +
 tests/99_extras.py      | 11 +++++++++++
 3 files changed, 13 insertions(+)

diff --git a/src/compiler/compiler.c b/src/compiler/compiler.c
index f72ccacb..b97ee107 100644
--- a/src/compiler/compiler.c
+++ b/src/compiler/compiler.c
@@ -2694,6 +2694,7 @@ static Error* compile_stmt(Compiler* self) {
                 vtdelete((Expr*)as_name);
                 if(!ok) return SyntaxError(self, "invalid syntax");
             } else {
+                // discard `__enter__()`'s return value
                 Ctx__emit_(ctx(), OP_POP_TOP, BC_NOARG, BC_KEEPLINE);
             }
             check(compile_block_body(self, compile_stmt));
diff --git a/src/interpreter/ceval.c b/src/interpreter/ceval.c
index d489cb55..58cb747b 100644
--- a/src/interpreter/ceval.c
+++ b/src/interpreter/ceval.c
@@ -1020,6 +1020,7 @@ FrameResult VM__run_top_frame(VM* self) {
                     goto __ERROR;
                 }
                 if(!py_vectorcall(0, 0)) goto __ERROR;
+                POP();
                 DISPATCH();
             }
             ///////////
diff --git a/tests/99_extras.py b/tests/99_extras.py
index 33d477c6..31612ac5 100644
--- a/tests/99_extras.py
+++ b/tests/99_extras.py
@@ -92,3 +92,14 @@ class fixed(TValue[int]):
         return super().__new__(cls, int(value))
     
 assert fixed('123').value == 123
+
+# context bug
+class Context:
+  def __enter__(self):
+    return 1
+  def __exit__(self, *_):
+    pass
+
+for _ in range(5):
+    with Context() as x:
+        assert x == 1