allow exception raise inside atexit callback

lib/utils/pyexec.c

@@ -44,6 +44,10 @@
 #include "lib/utils/pyexec.h"
 #include "genhdr/mpversion.h"
 
+#if CIRCUITPY_ATEXIT
+#include "shared-module/atexit/__init__.h"
+#endif
+
 pyexec_mode_kind_t pyexec_mode_kind = PYEXEC_MODE_FRIENDLY_REPL;
 int pyexec_system_exit = 0;
 
@@ -58,6 +62,7 @@ STATIC bool repl_display_debugging_info = 0;
 #define EXEC_FLAG_SOURCE_IS_VSTR (16)
 #define EXEC_FLAG_SOURCE_IS_FILENAME (32)
 #define EXEC_FLAG_SOURCE_IS_READER (64)
+#define EXEC_FLAG_SOURCE_IS_ATEXIT (128)
 
 // parses, compiles and executes the code in the lexer
 // frees the lexer before returning
@@ -81,6 +86,10 @@ STATIC int parse_compile_execute(const void *source, mp_parse_input_kind_t input
     nlr.ret_val = NULL;
     if (nlr_push(&nlr) == 0) {
         mp_obj_t module_fun;
+        #if CIRCUITPY_ATEXIT
+        if (!(exec_flags & EXEC_FLAG_SOURCE_IS_ATEXIT))
+        #endif
+        {
             #if MICROPY_MODULE_FROZEN_MPY
             if (exec_flags & EXEC_FLAG_SOURCE_IS_RAW_CODE) {
                 // source is a raw_code object, create the function
@@ -120,13 +129,22 @@ STATIC int parse_compile_execute(const void *source, mp_parse_input_kind_t input
                 module_fun = make_obj_long_lived(module_fun, 6);
                 gc_collect();
             }
+        }
 
         // execute code
         mp_hal_set_interrupt_char(CHAR_CTRL_C); // allow ctrl-C to interrupt us
         #if MICROPY_REPL_INFO
         start = mp_hal_ticks_ms();
         #endif
+        #if CIRCUITPY_ATEXIT
+        if (exec_flags & EXEC_FLAG_SOURCE_IS_ATEXIT) {
+            atexit_callback_t *callback = (atexit_callback_t *)source;
+            mp_call_function_n_kw(callback->func, callback->n_pos, callback->n_kw, callback->args);
+        } else
+        #endif
+        {
             mp_call_function_0(module_fun);
+        }
         mp_hal_set_interrupt_char(-1); // disable interrupt
         mp_handle_pending(true); // handle any pending exceptions (and any callbacks)
         nlr_pop();
@@ -741,6 +759,12 @@ int pyexec_frozen_module(const char *name, pyexec_result_t *result) {
 }
 #endif
 
+#if CIRCUITPY_ATEXIT
+int pyexec_exit_handler(const void *source, pyexec_result_t *result) {
+    return parse_compile_execute(source, MP_PARSE_FILE_INPUT, EXEC_FLAG_SOURCE_IS_ATEXIT, result);
+}
+#endif
+
 #if MICROPY_REPL_INFO
 mp_obj_t pyb_set_repl_info(mp_obj_t o_value) {
     repl_display_debugging_info = mp_obj_get_int(o_value);

lib/utils/pyexec.h

@@ -59,6 +59,10 @@ void pyexec_event_repl_init(void);
 int pyexec_event_repl_process_char(int c);
 extern uint8_t pyexec_repl_active;
 
+#if CIRCUITPY_ATEXIT
+int pyexec_exit_handler(const void *source, pyexec_result_t *result);
+#endif
+
 #if MICROPY_REPL_INFO
 mp_obj_t pyb_set_repl_info(mp_obj_t o_value);
 MP_DECLARE_CONST_FUN_OBJ_1(pyb_set_repl_info_obj);

main.c

@@ -218,7 +218,7 @@ STATIC bool maybe_run_list(const char * const * filenames, pyexec_result_t* exec
     mp_hal_stdout_tx_str(decompressed);
     pyexec_file(filename, exec_result);
     #if CIRCUITPY_ATEXIT
-    shared_module_atexit_execute();
+    shared_module_atexit_execute(exec_result);
     #endif
     return true;
 }
@@ -770,7 +770,11 @@ STATIC int run_repl(void) {
         exit_code = pyexec_friendly_repl();
     }
     #if CIRCUITPY_ATEXIT
-    shared_module_atexit_execute();
+    pyexec_result_t result;
+    shared_module_atexit_execute(&result);
+    if (result.return_code == PYEXEC_DEEP_SLEEP) {
+        exit_code = PYEXEC_DEEP_SLEEP;
+    }
     #endif
     cleanup_after_vm(heap, MP_OBJ_SENTINEL);
     #if CIRCUITPY_STATUS_LED

shared-module/atexit/__init__.c

@@ -28,11 +28,6 @@
 #include "py/runtime.h"
 #include "shared-module/atexit/__init__.h"
 
-typedef struct _atexit_callback_t {
-    size_t n_pos, n_kw;
-    mp_obj_t func, *args;
-} atexit_callback_t;
-
 static size_t callback_len = 0;
 static atexit_callback_t *callback = NULL;
 
@@ -79,11 +74,18 @@ void shared_module_atexit_unregister(const mp_obj_t *func) {
     }
 }
 
-void shared_module_atexit_execute(void) {
+void shared_module_atexit_execute(pyexec_result_t *result) {
     if (callback) {
         for (size_t i = callback_len; i-- > 0;) {
             if (callback[i].func != mp_const_none) {
-                mp_call_function_n_kw(callback[i].func, callback[i].n_pos, callback[i].n_kw, callback[i].args);
+                if (result != NULL) {
+                    pyexec_result_t res;
+                    if (pyexec_exit_handler(&callback[i], &res) == PYEXEC_DEEP_SLEEP) {
+                        *result = res;
+                    }
+                } else {
+                    pyexec_exit_handler(&callback[i], NULL);
+                }
             }
         }
     }

shared-module/atexit/__init__.h

@@ -28,6 +28,12 @@
 #define MICROPY_INCLUDED_SHARED_MODULE_ATEXIT___INIT___H
 
 #include "py/obj.h"
+#include "lib/utils/pyexec.h"
+
+typedef struct _atexit_callback_t {
+    size_t n_pos, n_kw;
+    mp_obj_t func, *args;
+} atexit_callback_t;
 
 extern void atexit_reset(void);
 extern void atexit_gc_collect(void);
@@ -35,6 +41,6 @@ extern void atexit_gc_collect(void);
 extern void shared_module_atexit_register(mp_obj_t *func,
     size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args);
 extern void shared_module_atexit_unregister(const mp_obj_t *func);
-extern void shared_module_atexit_execute(void);
+extern void shared_module_atexit_execute(pyexec_result_t *result);
 
 #endif  // MICROPY_INCLUDED_SHARED_MODULE_ATEXIT___INIT___H