Compare commits

...

3 Commits

Author SHA1 Message Date
Kanika Kapoor
ceaa4dc076
Merge c048ec9faf4bfe02da004dce69471897933c3617 into 5d30381e1915c1f3a148ef39e81002984d3951d8 2026-04-18 15:01:37 +00:00
blueloveTH
5d30381e19 fix #498 2026-04-17 11:10:55 +08:00
Kanika Kapoor
c048ec9faf Fix context manager __exit__ not being called on exception (#395)
Problem: When an exception occurs in a WITH block, __exit__ was not called,
preventing proper cleanup of context managers.

Solution:
1. Wrap WITH block body in try-except structure
2. On normal exit: call __exit__(None, None, None)
3. On exception: call __exit__ with exception info before re-raising

Changes:
- compiler.c: Wrap WITH body in try-except, ensure __exit__ called in both paths
- ceval.c: Update OP_WITH_EXIT to accept three arguments (exc_type, exc_val, exc_tb)
- tests/520_context.py: Add test to verify __exit__ called on exceptions
2025-12-27 01:12:15 +05:30
5 changed files with 95 additions and 14 deletions

View File

@ -2847,6 +2847,8 @@ static Error* compile_stmt(Compiler* self) {
case TK_WITH: {
check(EXPR(self)); // [ <expr> ]
Ctx__s_emit_top(ctx());
// Save context manager for later __exit__ call
Ctx__emit_(ctx(), OP_DUP_TOP, BC_NOARG, prev()->line);
Ctx__enter_block(ctx(), CodeBlockType_WITH);
NameExpr* as_name = NULL;
if(match(TK_AS)) {
@ -2855,17 +2857,33 @@ static Error* compile_stmt(Compiler* self) {
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 {
// discard `__enter__()`'s return value
Ctx__emit_(ctx(), OP_POP_TOP, BC_NOARG, BC_KEEPLINE);
}
// Wrap body in try-except to ensure __exit__ is called even on exception
Ctx__enter_block(ctx(), CodeBlockType_TRY);
Ctx__emit_(ctx(), OP_BEGIN_TRY, BC_NOARG, prev()->line);
check(compile_block_body(self));
Ctx__emit_(ctx(), OP_END_TRY, BC_NOARG, BC_KEEPLINE);
// Normal exit: call __exit__(None, None, None)
Ctx__emit_(ctx(), OP_LOAD_NONE, BC_NOARG, prev()->line);
Ctx__emit_(ctx(), OP_LOAD_NONE, BC_NOARG, prev()->line);
Ctx__emit_(ctx(), OP_LOAD_NONE, BC_NOARG, prev()->line);
Ctx__emit_(ctx(), OP_WITH_EXIT, BC_NOARG, prev()->line);
int jump_patch = Ctx__emit_(ctx(), OP_JUMP_FORWARD, BC_NOARG, BC_KEEPLINE);
Ctx__exit_block(ctx());
// Exception handler: call __exit__ with exception info, then re-raise
Ctx__emit_(ctx(), OP_PUSH_EXCEPTION, BC_NOARG, BC_KEEPLINE);
Ctx__emit_(ctx(), OP_LOAD_NONE, BC_NOARG, BC_KEEPLINE); // exc_type
Ctx__emit_(ctx(), OP_ROT_TWO, BC_NOARG, BC_KEEPLINE); // reorder: [cm, None, exc]
Ctx__emit_(ctx(), OP_LOAD_NONE, BC_NOARG, BC_KEEPLINE); // exc_tb
Ctx__emit_(ctx(), OP_WITH_EXIT, BC_NOARG, prev()->line);
Ctx__emit_(ctx(), OP_RE_RAISE, BC_NOARG, BC_KEEPLINE);
Ctx__patch_jump(ctx(), jump_patch);
Ctx__exit_block(ctx());
} break;
/*************************************************/

View File

@ -1122,14 +1122,35 @@ __NEXT_STEP:
DISPATCH();
}
case OP_WITH_EXIT: {
// [expr]
py_push(TOP());
// Stack: [cm, exc_type, exc_val, exc_tb]
// Call cm.__exit__(exc_type, exc_val, exc_tb)
py_Ref exc_tb = TOP();
py_Ref exc_val = SECOND();
py_Ref exc_type = THIRD();
py_Ref cm = FOURTH();
// Save all values from stack
py_TValue saved_cm = *cm;
py_TValue saved_exc_type = *exc_type;
py_TValue saved_exc_val = *exc_val;
py_TValue saved_exc_tb = *exc_tb;
self->stack.sp -= 4;
// Push cm and get __exit__ method
py_push(&saved_cm);
if(!py_pushmethod(__exit__)) {
TypeError("'%t' object does not support the context manager protocol", TOP()->type);
TypeError("'%t' object does not support the context manager protocol", saved_cm.type);
goto __ERROR;
}
if(!py_vectorcall(0, 0)) goto __ERROR;
POP();
// Push arguments: exc_type, exc_val, exc_tb
PUSH(&saved_exc_type);
PUSH(&saved_exc_val);
PUSH(&saved_exc_tb);
// Call __exit__(exc_type, exc_val, exc_tb)
if(!py_vectorcall(3, 0)) goto __ERROR;
py_pop(); // discard return value
DISPATCH();
}
///////////

View File

@ -136,6 +136,7 @@ static bool io_FileIO__exit__(int argc, py_Ref argv) {
static bool io_FileIO_read(int argc, py_Ref argv) {
io_FileIO* ud = py_touserdata(py_arg(0));
if(ud->file == NULL) return ValueError("I/O operation on closed file");
bool is_binary = ud->mode[strlen(ud->mode) - 1] == 'b';
int size;
if(argc == 1) {
@ -146,24 +147,37 @@ static bool io_FileIO_read(int argc, py_Ref argv) {
} else if(argc == 2) {
PY_CHECK_ARG_TYPE(1, tp_int);
size = py_toint(py_arg(1));
if (size < 0) {
long current = ftell(ud->file);
fseek(ud->file, 0, SEEK_END);
size = ftell(ud->file);
fseek(ud->file, current, SEEK_SET);
}
} else {
return TypeError("read() takes at most 2 arguments (%d given)", argc);
}
if(is_binary) {
void* dst = py_newbytes(py_retval(), size);
int actual_size = fread(dst, 1, size, ud->file);
py_bytes_resize(py_retval(), actual_size);
if(size > 0) {
int actual_size = fread(dst, 1, size, ud->file);
py_bytes_resize(py_retval(), actual_size);
}
} else {
void* dst = PK_MALLOC(size);
int actual_size = fread(dst, 1, size, ud->file);
py_newstrv(py_retval(), (c11_sv){dst, actual_size});
PK_FREE(dst);
if(size > 0) {
void* dst = PK_MALLOC(size);
int actual_size = fread(dst, 1, size, ud->file);
py_newstrv(py_retval(), (c11_sv){dst, actual_size});
PK_FREE(dst);
} else {
py_newstr(py_retval(), "");
}
}
return true;
}
static bool io_FileIO_tell(int argc, py_Ref argv) {
io_FileIO* ud = py_touserdata(py_arg(0));
if(ud->file == NULL) return ValueError("I/O operation on closed file");
py_newint(py_retval(), ftell(ud->file));
return true;
}
@ -173,6 +187,7 @@ static bool io_FileIO_seek(int argc, py_Ref argv) {
PY_CHECK_ARG_TYPE(1, tp_int);
PY_CHECK_ARG_TYPE(2, tp_int);
io_FileIO* ud = py_touserdata(py_arg(0));
if(ud->file == NULL) return ValueError("I/O operation on closed file");
long cookie = py_toint(py_arg(1));
int whence = py_toint(py_arg(2));
py_newint(py_retval(), fseek(ud->file, cookie, whence));
@ -193,6 +208,7 @@ static bool io_FileIO_close(int argc, py_Ref argv) {
static bool io_FileIO_write(int argc, py_Ref argv) {
PY_CHECK_ARGC(2);
io_FileIO* ud = py_touserdata(py_arg(0));
if(ud->file == NULL) return ValueError("I/O operation on closed file");
size_t written_size;
if(ud->mode[strlen(ud->mode) - 1] == 'b') {
PY_CHECK_ARG_TYPE(1, tp_bytes);
@ -211,6 +227,7 @@ static bool io_FileIO_write(int argc, py_Ref argv) {
static bool io_FileIO_flush(int argc, py_Ref argv) {
PY_CHECK_ARGC(1);
io_FileIO* ud = py_touserdata(py_arg(0));
if(ud->file == NULL) return ValueError("I/O operation on closed file");
fflush(ud->file);
py_newnone(py_retval());
return true;

View File

@ -27,4 +27,29 @@ assert path == ['enter', 'in', 'exit']
path.clear()
# Test that __exit__ is called even when an exception occurs
class B:
def __init__(self):
self.path = []
def __enter__(self):
path.append('enter')
return self
def __exit__(self, exc_type, exc_val, exc_tb):
path.append('exit')
if exc_type is not None:
path.append('exception')
return False # propagate exception
try:
with B():
path.append('before_raise')
raise ValueError('test')
path.append('after_raise') # should not be reached
except ValueError:
pass
assert path == ['enter', 'before_raise', 'exit', 'exception'], f"Expected ['enter', 'before_raise', 'exit', 'exception'], got {path}"

View File

@ -38,7 +38,7 @@ with open('123.txt', 'rt') as f:
assert f.read() == ''
f.seek(3, io.SEEK_SET)
assert f.tell() == 3
assert f.read() == '456'
assert f.read(-1) == '456'
assert f.tell() == 6
with open('123.txt', 'a') as f: