From 471d4ae578545fdc1db278657c09b6e2db295290 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@foxmail.com>
Date: Thu, 8 Aug 2024 12:24:17 +0800
Subject: [PATCH] ...

---
 src/compiler/compiler.c                | 45 +++++++++++++-------------
 src/interpreter/ceval.c                | 22 +++++++++++++
 tests/{67_context.py => 52_context.py} |  0
 3 files changed, 45 insertions(+), 22 deletions(-)
 rename tests/{67_context.py => 52_context.py} (100%)

diff --git a/src/compiler/compiler.c b/src/compiler/compiler.c
index f14209f1..b127a141 100644
--- a/src/compiler/compiler.c
+++ b/src/compiler/compiler.c
@@ -2664,28 +2664,29 @@ static Error* compile_stmt(Compiler* self) {
             Ctx__s_pop(ctx());
             consume_end_stmt();
         } break;
-        // case TK_WITH: {
-        //     check(EXPR(self));  // [ <expr> ]
-        //     Ctx__s_emit_top(ctx());
-        //     Ctx__enter_block(ctx(), CodeBlockType_CONTEXT_MANAGER);
-        //     Expr* as_name = nullptr;
-        //     if(match(TK_AS)) {
-        //         consume(TK_ID);
-        //         as_name = make_expr<NameExpr>(prev().str(), name_scope());
-        //     }
-        //     Ctx__emit_(ctx(), OP_WITH_ENTER, BC_NOARG, prev().line);
-        //     // [ <expr> <expr>.__enter__() ]
-        //     if(as_name) {
-        //         bool ok = as_name->emit_store(ctx());
-        //         delete_expr(as_name);
-        //         if(!ok) return SyntaxError(self, );
-        //     } else {
-        //         Ctx__emit_(ctx(), OP_POP_TOP, BC_NOARG, BC_KEEPLINE);
-        //     }
-        //     check(compile_block_body());
-        //     Ctx__emit_(ctx(), OP_WITH_EXIT, BC_NOARG, prev().line);
-        //     Ctx__exit_block(ctx());
-        // } break;
+        case TK_WITH: {
+            check(EXPR(self));  // [ <expr> ]
+            Ctx__s_emit_top(ctx());
+            Ctx__enter_block(ctx(), CodeBlockType_CONTEXT_MANAGER);
+            NameExpr* as_name = NULL;
+            if(match(TK_AS)) {
+                consume(TK_ID);
+                py_Name name = py_namev(Token__sv(prev()));
+                as_name = NameExpr__new(prev()->line, name, name_scope(self));
+            }
+            Ctx__emit_(ctx(), OP_WITH_ENTER, BC_NOARG, prev()->line);
+            // [ <expr> <expr>.__enter__() ]
+            if(as_name) {
+                bool ok = vtemit_store((Expr*)as_name, ctx());
+                vtdelete((Expr*)as_name);
+                if(!ok) return SyntaxError(self, "invalid syntax");
+            } else {
+                Ctx__emit_(ctx(), OP_POP_TOP, BC_NOARG, BC_KEEPLINE);
+            }
+            check(compile_block_body(self, compile_stmt));
+            Ctx__emit_(ctx(), OP_WITH_EXIT, BC_NOARG, prev()->line);
+            Ctx__exit_block(ctx());
+        } break;
         /*************************************************/
         case TK_EQ: {
             consume(TK_ID);
diff --git a/src/interpreter/ceval.c b/src/interpreter/ceval.c
index 3dc4e1f0..8970309b 100644
--- a/src/interpreter/ceval.c
+++ b/src/interpreter/ceval.c
@@ -898,6 +898,28 @@ FrameResult VM__run_top_frame(VM* self) {
                 DISPATCH();
             }
             ///////////
+            case OP_WITH_ENTER: {
+                // [expr]
+                py_push(TOP());
+                if(!py_pushmethod(__enter__)){
+                    TypeError("'%t' object does not support the context manager protocol", TOP()->type);
+                    goto __ERROR;
+                }
+                if(!py_vectorcall(0, 0)) goto __ERROR;
+                PUSH(py_retval());
+                DISPATCH();
+            }
+            case OP_WITH_EXIT: {
+                // [expr]
+                py_push(TOP());
+                if(!py_pushmethod(__exit__)){
+                    TypeError("'%t' object does not support the context manager protocol", TOP()->type);
+                    goto __ERROR;
+                }
+                if(!py_vectorcall(0, 0)) goto __ERROR;
+                DISPATCH();
+            }
+            ///////////
             case OP_TRY_ENTER: {
                 Frame__set_unwind_target(frame, SP());
                 DISPATCH();
diff --git a/tests/67_context.py b/tests/52_context.py
similarity index 100%
rename from tests/67_context.py
rename to tests/52_context.py