Skip to content

Commit 4bf3f2d

Browse files
dpgeorgedpgeorge
committed
py: Fix with+for+return bug by popping for-iter when unwinding exc stack.
Addresses issue adafruit#1182.
1 parent 556c8a9 commit 4bf3f2d

3 files changed

Lines changed: 110 additions & 9 deletions

File tree

‎py/vm.c‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -637,10 +637,14 @@ unwind_jump:;
637637
unum -= 1;
638638
assert(exc_sp >= exc_stack);
639639
if (MP_TAGPTR_TAG1(exc_sp->val_sp)) {
640+
// Getting here the stack looks like:
641+
// (..., X, dest_ip)
642+
// where X is pointed to by exc_sp->val_sp and in the case
643+
// of a "with" block contains the context manager info.
640644
// We're going to run "finally" code as a coroutine
641645
// (not calling it recursively). Set up a sentinel
642646
// on a stack so it can return back to us when it is
643-
// done (when END_FINALLY reached).
647+
// done (when WITH_CLEANUP or END_FINALLY reached).
644648
PUSH((void*)unum); // push number of exception handlers left to unwind
645649
PUSH(MP_OBJ_NEW_SMALL_INT(UNWIND_JUMP)); // push sentinel
646650
ip = exc_sp->handler; // get exception handler byte code address
@@ -1016,15 +1020,24 @@ unwind_jump:;
10161020
unwind_return:
10171021
while (exc_sp >= exc_stack) {
10181022
if (MP_TAGPTR_TAG1(exc_sp->val_sp)) {
1023+
// Getting here the stack looks like:
1024+
// (..., X, [iter0, iter1, ...,] ret_val)
1025+
// where X is pointed to by exc_sp->val_sp and in the case
1026+
// of a "with" block contains the context manager info.
1027+
// There may be 0 or more for-iterators between X and the
1028+
// return value, and these must be removed before control can
1029+
// pass to the finally code. We simply copy the ret_value down
1030+
// over these iterators, if they exist. If they don't then the
1031+
// following is a null operation.
1032+
mp_obj_t *finally_sp = MP_TAGPTR_PTR(exc_sp->val_sp);
1033+
finally_sp[1] = sp[0];
1034+
sp = &finally_sp[1];
10191035
// We're going to run "finally" code as a coroutine
10201036
// (not calling it recursively). Set up a sentinel
10211037
// on a stack so it can return back to us when it is
1022-
// done (when END_FINALLY reached).
1038+
// done (when WITH_CLEANUP or END_FINALLY reached).
10231039
PUSH(MP_OBJ_NEW_SMALL_INT(UNWIND_RETURN));
10241040
ip = exc_sp->handler;
1025-
// We don't need to do anything with sp, finally is just
1026-
// syntactic sugar for sequential execution??
1027-
// sp =
10281041
exc_sp--;
10291042
goto dispatch_loop;
10301043
}

‎tests/basics/try_finally_return.py‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,3 +21,52 @@ def func3():
2121
print("finally 3")
2222

2323
print(func3())
24+
25+
# for loop within try-finally
26+
def f():
27+
try:
28+
for i in [1, 2]:
29+
return i
30+
finally:
31+
print('finally')
32+
print(f())
33+
34+
# multiple for loops within try-finally
35+
def f():
36+
try:
37+
for i in [1, 2]:
38+
for j in [3, 4]:
39+
return (i, j)
40+
finally:
41+
print('finally')
42+
print(f())
43+
44+
# multiple for loops and nested try-finally's
45+
def f():
46+
try:
47+
for i in [1, 2]:
48+
for j in [3, 4]:
49+
try:
50+
for k in [5, 6]:
51+
for l in [7, 8]:
52+
return (i, j, k, l)
53+
finally:
54+
print('finally 2')
55+
finally:
56+
print('finally 1')
57+
print(f())
58+
59+
# multiple for loops that are optimised, and nested try-finally's
60+
def f():
61+
try:
62+
for i in range(1, 3):
63+
for j in range(3, 5):
64+
try:
65+
for k in range(5, 7):
66+
for l in range(7, 9):
67+
return (i, j, k, l)
68+
finally:
69+
print('finally 2')
70+
finally:
71+
print('finally 1')
72+
print(f())

‎tests/basics/with_return.py‎

Lines changed: 43 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,53 @@
11
class CtxMgr:
2+
def __init__(self, id):
3+
self.id = id
24

35
def __enter__(self):
4-
print("__enter__")
6+
print("__enter__", self.id)
57
return self
68

79
def __exit__(self, a, b, c):
8-
print("__exit__", repr(a), repr(b))
10+
print("__exit__", self.id, repr(a), repr(b))
911

12+
# simple case
1013
def foo():
11-
with CtxMgr():
14+
with CtxMgr(1):
1215
return 4
13-
1416
print(foo())
17+
18+
# for loop within with (iterator needs removing upon return)
19+
def f():
20+
with CtxMgr(1):
21+
for i in [1, 2]:
22+
return i
23+
print(f())
24+
25+
# multiple for loops within with
26+
def f():
27+
with CtxMgr(1):
28+
for i in [1, 2]:
29+
for j in [3, 4]:
30+
return (i, j)
31+
print(f())
32+
33+
# multiple for loops within nested withs
34+
def f():
35+
with CtxMgr(1):
36+
for i in [1, 2]:
37+
for j in [3, 4]:
38+
with CtxMgr(2):
39+
for k in [5, 6]:
40+
for l in [7, 8]:
41+
return (i, j, k, l)
42+
print(f())
43+
44+
# multiple for loops that are optimised, and nested withs
45+
def f():
46+
with CtxMgr(1):
47+
for i in range(1, 3):
48+
for j in range(3, 5):
49+
with CtxMgr(2):
50+
for k in range(5, 7):
51+
for l in range(7, 9):
52+
return (i, j, k, l)
53+
print(f())

0 commit comments

Comments
 (0)