From dd443bacb857ae736bc51d81585e1063cd0a03f2 Mon Sep 17 00:00:00 2001
From: Jeff Epler <jepler@gmail.com>
Date: Mon, 17 Oct 2022 09:08:38 -0500
Subject: [PATCH] Chain exceptions while unwinding

---
 py/vm.c                                           | 11 ++++++++++-
 tests/basics/exception_chain.py                   | 10 +++++++++-
 tests/circuitpython/traceback_test_chained.py     |  8 ++++++++
 tests/circuitpython/traceback_test_chained.py.exp | 12 ++++++++++++
 4 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/py/vm.c b/py/vm.c
index 6d9b7f7a0b..9e9dbcaa8d 100644
--- a/py/vm.c
+++ b/py/vm.c
@@ -1456,10 +1456,19 @@ unwind_loop:
                 // catch exception and pass to byte code
                 code_state->ip = exc_sp->handler;
                 mp_obj_t *sp = MP_TAGPTR_PTR(exc_sp->val_sp);
+                #if MICROPY_CPYTHON_EXCEPTION_CHAIN
+                mp_obj_t active_exception = get_active_exception(exc_sp, exc_stack);
+                #endif
                 // save this exception in the stack so it can be used in a reraise, if needed
                 exc_sp->prev_exc = nlr.ret_val;
+                mp_obj_t obj = MP_OBJ_FROM_PTR(nlr.ret_val);
+                #if MICROPY_CPYTHON_EXCEPTION_CHAIN
+                if (active_exception != MP_OBJ_NULL) {
+                    mp_store_attr(obj, MP_QSTR___context__, active_exception);
+                }
+                #endif
                 // push exception object so it can be handled by bytecode
-                PUSH(MP_OBJ_FROM_PTR(nlr.ret_val));
+                PUSH(obj);
                 code_state->sp = sp;
 
             #if MICROPY_STACKLESS
diff --git a/tests/basics/exception_chain.py b/tests/basics/exception_chain.py
index ef506c7cb0..b84ff19dd6 100644
--- a/tests/basics/exception_chain.py
+++ b/tests/basics/exception_chain.py
@@ -5,7 +5,7 @@ except AttributeError:
     raise SystemExit
 
 def print_exc_info(e):
-    print("exception", type(e), repr(e))
+    print("exception", type(e), e.args)
     print("context", type(e.__context__), e.__suppress_context__)
     print("cause", type(e.__cause__))
 
@@ -44,3 +44,11 @@ try:
         raise RuntimeError() from None
 except Exception as e:
     print_exc_info(e)
+
+try:
+    try:
+        raise RuntimeError()
+    except Exception as inner:
+        1/0
+except Exception as e:
+    print_exc_info(e)
diff --git a/tests/circuitpython/traceback_test_chained.py b/tests/circuitpython/traceback_test_chained.py
index a0acc876cc..83d0c03466 100644
--- a/tests/circuitpython/traceback_test_chained.py
+++ b/tests/circuitpython/traceback_test_chained.py
@@ -53,3 +53,11 @@ try:
         raise RuntimeError() from None
 except Exception as e:
     print_exc_info(e)
+
+try:
+    try:
+        raise RuntimeError()
+    except Exception as inner:
+        1 / 0
+except Exception as e:
+    print_exc_info(e)
diff --git a/tests/circuitpython/traceback_test_chained.py.exp b/tests/circuitpython/traceback_test_chained.py.exp
index bc26c44a3e..c874ff707f 100644
--- a/tests/circuitpython/traceback_test_chained.py.exp
+++ b/tests/circuitpython/traceback_test_chained.py.exp
@@ -41,3 +41,15 @@ Traceback (most recent call last):
 RuntimeError: 
 ------------------------------------------------------------------------
 
+------------------------------------------------------------------------
+Traceback (most recent call last):
+  File "circuitpython/traceback_test_chained.py", line 59, in <module>
+RuntimeError: 
+
+During handling of the above exception, another exception occurred:
+
+Traceback (most recent call last):
+  File "circuitpython/traceback_test_chained.py", line 61, in <module>
+ZeroDivisionError: division by zero
+------------------------------------------------------------------------
+