Compare commits

...

4 Commits

Author SHA1 Message Date
Kanika Kapoor
ff61ebf51b
Merge c048ec9faf4bfe02da004dce69471897933c3617 into 979addecf9cb180c59c16ac605ba75441ece5af6 2026-01-15 12:43:39 +05:30
blueloveTH
979addecf9 improve extend and choice 2026-01-14 16:49:02 +08:00
blueloveTH
9ef38d605b fix bool-int ops 2026-01-14 16:14:08 +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
9 changed files with 184 additions and 15 deletions

View File

@ -567,6 +567,66 @@ static bool bool__invert__(int argc, py_Ref argv) {
return true; return true;
} }
static bool bool_try_cast_i64(py_Ref arg, py_i64* out) {
if (arg->type == tp_int) {
*out = py_toint(arg);
return true;
} else if (arg->type == tp_bool) {
*out = py_tobool(arg);
return true;
} else {
return false;
}
}
static bool bool__add__(int argc, py_Ref argv) {
PY_CHECK_ARGC(2);
py_i64 lhs = py_tobool(py_arg(0));
py_i64 rhs;
if (bool_try_cast_i64(py_arg(1), &rhs)) {
py_newint(py_retval(), lhs + rhs);
} else {
py_newnotimplemented(py_retval());
}
return true;
}
static bool bool__sub__(int argc, py_Ref argv) {
PY_CHECK_ARGC(2);
py_i64 lhs = py_tobool(py_arg(0));
py_i64 rhs;
if (bool_try_cast_i64(py_arg(1), &rhs)) {
py_newint(py_retval(), lhs - rhs);
} else {
py_newnotimplemented(py_retval());
}
return true;
}
static bool bool__rsub__(int argc, py_Ref argv) {
PY_CHECK_ARGC(2);
py_i64 lhs = py_tobool(py_arg(0));
py_i64 rhs;
if (bool_try_cast_i64(py_arg(1), &rhs)) {
py_newint(py_retval(), rhs - lhs);
} else {
py_newnotimplemented(py_retval());
}
return true;
}
static bool bool__mul__(int argc, py_Ref argv) {
PY_CHECK_ARGC(2);
py_i64 lhs = py_tobool(py_arg(0));
py_i64 rhs;
if (bool_try_cast_i64(py_arg(1), &rhs)) {
py_newint(py_retval(), lhs * rhs);
} else {
py_newnotimplemented(py_retval());
}
return true;
}
void pk_number__register() { void pk_number__register() {
/****** tp_int & tp_float ******/ /****** tp_int & tp_float ******/
py_bindmagic(tp_int, __add__, int__add__); py_bindmagic(tp_int, __add__, int__add__);
@ -651,6 +711,12 @@ void pk_number__register() {
py_bindmagic(tp_bool, __or__, bool__or__); py_bindmagic(tp_bool, __or__, bool__or__);
py_bindmagic(tp_bool, __xor__, bool__xor__); py_bindmagic(tp_bool, __xor__, bool__xor__);
py_bindmagic(tp_bool, __invert__, bool__invert__); py_bindmagic(tp_bool, __invert__, bool__invert__);
py_bindmagic(tp_bool, __add__, bool__add__);
py_bindmagic(tp_bool, __sub__, bool__sub__);
py_bindmagic(tp_bool, __mul__, bool__mul__);
py_bindmagic(tp_bool, __radd__, bool__add__);
py_bindmagic(tp_bool, __rsub__, bool__rsub__);
py_bindmagic(tp_bool, __rmul__, bool__mul__);
} }
#undef DEF_NUM_BINARY_OP #undef DEF_NUM_BINARY_OP

View File

@ -2758,6 +2758,8 @@ static Error* compile_stmt(Compiler* self) {
case TK_WITH: { case TK_WITH: {
check(EXPR(self)); // [ <expr> ] check(EXPR(self)); // [ <expr> ]
Ctx__s_emit_top(ctx()); 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); Ctx__enter_block(ctx(), CodeBlockType_WITH);
NameExpr* as_name = NULL; NameExpr* as_name = NULL;
if(match(TK_AS)) { if(match(TK_AS)) {
@ -2766,17 +2768,33 @@ static Error* compile_stmt(Compiler* self) {
as_name = NameExpr__new(prev()->line, name, name_scope(self)); as_name = NameExpr__new(prev()->line, name, name_scope(self));
} }
Ctx__emit_(ctx(), OP_WITH_ENTER, BC_NOARG, prev()->line); Ctx__emit_(ctx(), OP_WITH_ENTER, BC_NOARG, prev()->line);
// [ <expr> <expr>.__enter__() ]
if(as_name) { if(as_name) {
bool ok = vtemit_store((Expr*)as_name, ctx()); bool ok = vtemit_store((Expr*)as_name, ctx());
vtdelete((Expr*)as_name); vtdelete((Expr*)as_name);
if(!ok) return SyntaxError(self, "invalid syntax"); if(!ok) return SyntaxError(self, "invalid syntax");
} else { } else {
// discard `__enter__()`'s return value
Ctx__emit_(ctx(), OP_POP_TOP, BC_NOARG, BC_KEEPLINE); 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)); 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); 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()); Ctx__exit_block(ctx());
} break; } break;
/*************************************************/ /*************************************************/

View File

@ -1122,14 +1122,35 @@ __NEXT_STEP:
DISPATCH(); DISPATCH();
} }
case OP_WITH_EXIT: { case OP_WITH_EXIT: {
// [expr] // Stack: [cm, exc_type, exc_val, exc_tb]
py_push(TOP()); // 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__)) { 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; 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(); DISPATCH();
} }
/////////// ///////////

View File

@ -218,12 +218,21 @@ static bool Random_randint(int argc, py_Ref argv) {
static bool Random_choice(int argc, py_Ref argv) { static bool Random_choice(int argc, py_Ref argv) {
PY_CHECK_ARGC(2); PY_CHECK_ARGC(2);
mt19937* ud = py_touserdata(py_arg(0)); mt19937* ud = py_touserdata(py_arg(0));
py_TValue* p; if (py_isstr(py_arg(1))) {
int length = pk_arrayview(py_arg(1), &p); c11_sv sv = py_tosv(py_arg(1));
if(length == -1) return TypeError("choice(): argument must be a list or tuple"); int length = c11_sv__u8_length(sv);
if(length == 0) return IndexError("cannot choose from an empty sequence"); if(length == 0) return IndexError("cannot choose from an empty sequence");
int index = mt19937__randint(ud, 0, length - 1); int index = mt19937__randint(ud, 0, length - 1);
py_assign(py_retval(), p + index); c11_sv ch = c11_sv__u8_getitem(sv, index);
py_newstrv(py_retval(), ch);
} else {
py_TValue* p;
int length = pk_arrayview(py_arg(1), &p);
if(length == -1) return TypeError("choice(): argument must be a list, tuple or str");
if(length == 0) return IndexError("cannot choose from an empty sequence");
int index = mt19937__randint(ud, 0, length - 1);
py_assign(py_retval(), p + index);
}
return true; return true;
} }

View File

@ -263,8 +263,22 @@ static bool list_extend(int argc, py_Ref argv) {
List* self = py_touserdata(py_arg(0)); List* self = py_touserdata(py_arg(0));
py_TValue* p; py_TValue* p;
int length = pk_arrayview(py_arg(1), &p); int length = pk_arrayview(py_arg(1), &p);
if(length == -1) return TypeError("extend() argument must be a list or tuple"); if(length >= 0) {
c11_vector__extend(self, p, length); c11_vector__extend(self, p, length);
} else {
// get iterator
if (!py_iter(py_arg(1))) return false;
py_StackRef tmp_iter = py_pushtmp();
py_assign(tmp_iter, py_retval());
while(true) {
int res = py_next(tmp_iter);
if (res == 0) break;
if (res == -1) return false;
assert(res == 1);
c11_vector__push(py_TValue, self, *py_retval());
}
py_pop();
}
py_newnone(py_retval()); py_newnone(py_retval());
return true; return true;
} }

View File

@ -31,3 +31,9 @@ assert NotImplemented is NotImplemented
assert True is True assert True is True
assert False is False assert False is False
assert True + 1 == 2
assert True - 1 == 0
assert True * 3 == 3
assert 1 + True == 2
assert 1 - True == 0
assert 3 * True == 3

View File

@ -170,6 +170,13 @@ b = a.copy(); del b[-1:]; assert b == [1, 2, 3]
b = a.copy(); del b[:]; assert b == [] b = a.copy(); del b[:]; assert b == []
assert a == [1, 2, 3, 4] assert a == [1, 2, 3, 4]
# test extend with iterable
c = [1]
c.extend('123')
assert c == [1, '1', '2', '3']
c.extend(range(1, 6))
assert c == [1, '1', '2', '3', 1, 2, 3, 4, 5]
# test cyclic reference # test cyclic reference
# a = [] # a = []
# a.append(0) # a.append(0)

View File

@ -27,4 +27,29 @@ assert path == ['enter', 'in', 'exit']
path.clear() 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

@ -20,6 +20,9 @@ for i in range(10):
for i in range(10): for i in range(10):
assert r.choice(tuple(a)) in a assert r.choice(tuple(a)) in a
for i in range(10):
assert r.choice('hello') in 'hello'
for i in range(10): for i in range(10):
assert r.randint(1, 1) == 1 assert r.randint(1, 1) == 1