Skip to content

Commit 21a07dc

Browse files
dpgeorgedpgeorge
committed
Merge pull request adafruit#389 from pfalcon/with-statement
With statement implementation
2 parents b04be05 + e7286ef commit 21a07dc

7 files changed

Lines changed: 186 additions & 2 deletions

File tree

‎py/compile.c‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1798,13 +1798,15 @@ void compile_with_stmt_helper(compiler_t *comp, int n, mp_parse_node_t *nodes, m
17981798
EMIT_ARG(setup_with, l_end);
17991799
EMIT(pop_top);
18001800
}
1801+
compile_increase_except_level(comp);
18011802
// compile additional pre-bits and the body
18021803
compile_with_stmt_helper(comp, n - 1, nodes + 1, body);
18031804
// finish this with block
18041805
EMIT(pop_block);
18051806
EMIT_ARG(load_const_tok, MP_TOKEN_KW_NONE);
18061807
EMIT_ARG(label_assign, l_end);
18071808
EMIT(with_cleanup);
1809+
compile_decrease_except_level(comp);
18081810
EMIT(end_finally);
18091811
}
18101812
}

‎py/qstrdefs.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ Q(__qualname__)
1616
Q(__repl_print__)
1717

1818
Q(__bool__)
19+
Q(__enter__)
20+
Q(__exit__)
1921
Q(__len__)
2022
Q(__iter__)
2123
Q(__getitem__)

‎py/vm.c‎

Lines changed: 69 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -381,6 +381,73 @@ mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **i
381381
break;
382382
*/
383383

384+
case MP_BC_SETUP_WITH: {
385+
obj1 = TOP();
386+
SET_TOP(rt_load_attr(obj1, MP_QSTR___exit__));
387+
mp_obj_t dest[2];
388+
rt_load_method(obj1, MP_QSTR___enter__, dest);
389+
obj2 = rt_call_method_n_kw(0, 0, dest);
390+
SETUP_BLOCK();
391+
PUSH(obj2);
392+
break;
393+
}
394+
395+
case MP_BC_WITH_CLEANUP: {
396+
static const mp_obj_t no_exc[] = {mp_const_none, mp_const_none, mp_const_none};
397+
if (TOP() == mp_const_none) {
398+
sp--;
399+
obj1 = TOP();
400+
SET_TOP(mp_const_none);
401+
obj2 = rt_call_function_n_kw(obj1, 3, 0, no_exc);
402+
} else if (MP_OBJ_IS_SMALL_INT(TOP())) {
403+
mp_obj_t cause = POP();
404+
switch (MP_OBJ_SMALL_INT_VALUE(cause)) {
405+
case UNWIND_RETURN: {
406+
mp_obj_t retval = POP();
407+
obj2 = rt_call_function_n_kw(TOP(), 3, 0, no_exc);
408+
SET_TOP(retval);
409+
PUSH(cause);
410+
break;
411+
}
412+
case UNWIND_JUMP: {
413+
obj2 = rt_call_function_n_kw(sp[-2], 3, 0, no_exc);
414+
// Pop __exit__ boundmethod at sp[-2]
415+
sp[-2] = sp[-1];
416+
sp[-1] = sp[0];
417+
SET_TOP(cause);
418+
break;
419+
}
420+
default:
421+
assert(0);
422+
}
423+
} else if (mp_obj_is_exception_type(TOP())) {
424+
mp_obj_t args[3] = {sp[0], sp[-1], sp[-2]};
425+
obj2 = rt_call_function_n_kw(sp[-3], 3, 0, args);
426+
// Pop __exit__ boundmethod at sp[-3]
427+
// TODO: Once semantics is proven, optimize for case when obj2 == True
428+
sp[-3] = sp[-2];
429+
sp[-2] = sp[-1];
430+
sp[-1] = sp[0];
431+
sp--;
432+
if (rt_is_true(obj2)) {
433+
// This is what CPython does
434+
//PUSH(MP_OBJ_NEW_SMALL_INT(UNWIND_SILENCED));
435+
// But what we need to do is - pop exception from value stack...
436+
sp -= 3;
437+
// ... pop with exception handler, and signal END_FINALLY
438+
// to just execute finally handler normally (signalled by None
439+
// on value stack)
440+
assert(exc_sp >= exc_stack);
441+
assert(exc_sp->opcode == MP_BC_SETUP_WITH);
442+
exc_sp--;
443+
PUSH(mp_const_none);
444+
}
445+
} else {
446+
assert(0);
447+
}
448+
break;
449+
}
450+
384451
case MP_BC_UNWIND_JUMP:
385452
DECODE_SLABEL;
386453
PUSH((void*)(ip + unum)); // push destination ip for jump
@@ -390,7 +457,7 @@ mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **i
390457
while (unum > 0) {
391458
unum -= 1;
392459
assert(exc_sp >= exc_stack);
393-
if (exc_sp->opcode == MP_BC_SETUP_FINALLY) {
460+
if (exc_sp->opcode == MP_BC_SETUP_FINALLY || exc_sp->opcode == MP_BC_SETUP_WITH) {
394461
// We're going to run "finally" code as a coroutine
395462
// (not calling it recursively). Set up a sentinel
396463
// on a stack so it can return back to us when it is
@@ -604,7 +671,7 @@ mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **i
604671
case MP_BC_RETURN_VALUE:
605672
unwind_return:
606673
while (exc_sp >= exc_stack) {
607-
if (exc_sp->opcode == MP_BC_SETUP_FINALLY) {
674+
if (exc_sp->opcode == MP_BC_SETUP_FINALLY || exc_sp->opcode == MP_BC_SETUP_WITH) {
608675
// We're going to run "finally" code as a coroutine
609676
// (not calling it recursively). Set up a sentinel
610677
// on a stack so it can return back to us when it is

‎tests/basics/with-break.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
class CtxMgr:
2+
3+
def __enter__(self):
4+
print("__enter__")
5+
return self
6+
7+
def __exit__(self, a, b, c):
8+
print("__exit__", repr(a), repr(b))
9+
10+
for i in range(5):
11+
print(i)
12+
with CtxMgr():
13+
if i == 3:
14+
break

‎tests/basics/with-continue.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
class CtxMgr:
2+
3+
def __enter__(self):
4+
print("__enter__")
5+
return self
6+
7+
def __exit__(self, a, b, c):
8+
print("__exit__", repr(a), repr(b))
9+
10+
for i in range(5):
11+
print(i)
12+
with CtxMgr():
13+
if i == 3:
14+
continue

‎tests/basics/with-return.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
class CtxMgr:
2+
3+
def __enter__(self):
4+
print("__enter__")
5+
return self
6+
7+
def __exit__(self, a, b, c):
8+
print("__exit__", repr(a), repr(b))
9+
10+
def foo():
11+
with CtxMgr():
12+
return 4
13+
14+
print(foo())

‎tests/basics/with1.py‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
class CtxMgr:
2+
3+
def __enter__(self):
4+
print("__enter__")
5+
return self
6+
7+
def __exit__(self, a, b, c):
8+
print("__exit__", repr(a), repr(b))
9+
10+
11+
with CtxMgr() as a:
12+
print(isinstance(a, CtxMgr))
13+
14+
try:
15+
with CtxMgr() as a:
16+
raise ValueError
17+
except ValueError:
18+
print("ValueError")
19+
20+
21+
class CtxMgr2:
22+
23+
def __enter__(self):
24+
print("__enter__")
25+
return self
26+
27+
def __exit__(self, a, b, c):
28+
print("__exit__", repr(a), repr(b))
29+
return True
30+
31+
try:
32+
with CtxMgr2() as a:
33+
raise ValueError
34+
print("No ValueError2")
35+
except ValueError:
36+
print("ValueError2")
37+
38+
39+
# These recursive try-finally tests are attempt to get some interpretation
40+
# of last phrase in http://docs.python.org/3.4/library/dis.html#opcode-WITH_CLEANUP
41+
# "If the stack represents an exception, and the function call returns a ‘true’
42+
# value, this information is “zapped” and replaced with a single WHY_SILENCED
43+
# to prevent END_FINALLY from re-raising the exception. (But non-local gotos
44+
# will still be resumed.)"
45+
print("===")
46+
with CtxMgr2() as a:
47+
try:
48+
try:
49+
raise ValueError
50+
print("No ValueError3")
51+
finally:
52+
print("finally1")
53+
finally:
54+
print("finally2")
55+
56+
print("===")
57+
try:
58+
try:
59+
with CtxMgr2() as a:
60+
try:
61+
try:
62+
raise ValueError
63+
print("No ValueError3")
64+
finally:
65+
print("finally1")
66+
finally:
67+
print("finally2")
68+
finally:
69+
print("finally3")
70+
finally:
71+
print("finally4")

0 commit comments

Comments
 (0)