py/compile: Fix async-for/async-with to work with simpler exc on stack.

There is now just the exception instance on the stack when an exception is
raised, not the full (type, exc, traceback).

5 files changed, 33 insertions, 10 deletions

diff --git a/py/compile.c b/py/compile.c
index 2fae5c9f64..3039f98b5d 100644
--- a/py/compile.c
+++ b/py/compile.c
@@ -1710,14 +1710,12 @@ STATIC void compile_async_for_stmt(compiler_t *comp, mp_parse_node_struct_t *pns
     EMIT_LOAD_GLOBAL(MP_QSTR_StopAsyncIteration);
     EMIT_ARG(binary_op, MP_BINARY_OP_EXCEPTION_MATCH);
     EMIT_ARG(pop_jump_if, false, try_finally_label);
-    EMIT(pop_top);
-    EMIT(pop_top);
-    EMIT(pop_top);
+    EMIT(pop_top); // pop exception instance
     EMIT(pop_except);
     EMIT_ARG(jump, while_else_label);
 
     EMIT_ARG(label_assign, try_finally_label);
-    EMIT_ARG(adjust_stack_size, 3);
+    EMIT_ARG(adjust_stack_size, 1); // if we jump here, the exc is on the stack
     compile_decrease_except_level(comp);
     EMIT(end_finally);
     EMIT(end_except_handler);
@@ -1778,9 +1776,21 @@ STATIC void compile_async_with_stmt_helper(compiler_t *comp, int n, mp_parse_nod
 
         EMIT_ARG(label_assign, try_exception_label); // start of exception handler
         EMIT(start_except_handler);
-        EMIT(rot_three);
+
+        // at this point the stack contains: ..., __aexit__, self, exc
+        EMIT(dup_top);
+        #if MICROPY_CPYTHON_COMPAT
+        EMIT_ARG(load_attr, MP_QSTR___class__); // get type(exc)
+        #else
+        compile_load_id(comp, MP_QSTR_type);
+        EMIT(rot_two);
+        EMIT_ARG(call_function, 1, 0, 0); // get type(exc)
+        #endif
         EMIT(rot_two);
+        EMIT_ARG(load_const_tok, MP_TOKEN_KW_NONE); // dummy traceback value
+        // at this point the stack contains: ..., __aexit__, self, type(exc), exc, None
         EMIT_ARG(call_method, 3, 0, 0);
+
         compile_yield_from(comp);
         EMIT_ARG(pop_jump_if, true, no_reraise_label);
         EMIT_ARG(raise_varargs, 0);
@@ -1789,7 +1799,7 @@ STATIC void compile_async_with_stmt_helper(compiler_t *comp, int n, mp_parse_nod
         EMIT(pop_except);
         EMIT_ARG(jump, end_label);
 
-        EMIT_ARG(adjust_stack_size, 5);
+        EMIT_ARG(adjust_stack_size, 3); // adjust for __aexit__, self, exc
         compile_decrease_except_level(comp);
         EMIT(end_finally);
         EMIT(end_except_handler);
diff --git a/tests/basics/async_with.py b/tests/basics/async_with.py
index 742f9ba993..9eccfd816c 100644
--- a/tests/basics/async_with.py
+++ b/tests/basics/async_with.py
@@ -4,7 +4,7 @@ class AContext:
     async def __aenter__(self):
         print('enter')
     async def __aexit__(self, exc_type, exc, tb):
-        print('exit')
+        print('exit', exc_type, exc)
 
 async def f():
     async with AContext():
@@ -15,3 +15,13 @@ try:
     o.send(None)
 except StopIteration:
     print('finished')
+
+async def g():
+    async with AContext():
+        raise ValueError('error')
+
+o = g()
+try:
+    o.send(None)
+except ValueError:
+    print('ValueError')
diff --git a/tests/basics/async_with.py.exp b/tests/basics/async_with.py.exp
index 1e9176af7b..6072a3e0b3 100644
--- a/tests/basics/async_with.py.exp
+++ b/tests/basics/async_with.py.exp
@@ -1,4 +1,7 @@
 enter
 body
-exit
+exit None None
 finished
+enter
+exit <class 'ValueError'> error
+ValueError
diff --git a/tests/basics/async_with2.py b/tests/basics/async_with2.py
index 0ebec489fe..44421ae917 100644
--- a/tests/basics/async_with2.py
+++ b/tests/basics/async_with2.py
@@ -20,7 +20,7 @@ class AContext:
         print('enter')
         print('f returned:', await f(10))
     async def __aexit__(self, exc_type, exc, tb):
-        print('exit')
+        print('exit', exc_type, exc)
         print('f returned:', await f(20))
 
 async def coro():
diff --git a/tests/basics/async_with2.py.exp b/tests/basics/async_with2.py.exp
index dd5a1c549a..76b173b4c2 100644
--- a/tests/basics/async_with2.py.exp
+++ b/tests/basics/async_with2.py.exp
@@ -9,7 +9,7 @@ coro yielded: 31
 coro yielded: 32
 body f returned: 33
 body end
-exit
+exit None None
 f start: 20
 coro yielded: 21
 coro yielded: 22