py/modsys: Extend atexit to hold multiple functions.

pi-anl · andrewleech · commit a952fe11516b · 2024-10-25T08:02:31.000+11:00
Signed-off-by: Andrew Leech &lt;andrew@alelec.net&gt;
diff --git a/py/modsys.c b/py/modsys.c
@@ -208,13 +208,42 @@ static MP_DEFINE_CONST_FUN_OBJ_1(mp_sys_getsizeof_obj, mp_sys_getsizeof);
 #endif
 
 #if MICROPY_PY_SYS_ATEXIT
+typedef struct _m_atexit_node_t {
+    struct _m_atexit_node_t *prev;
+    struct _m_atexit_node_t *next;
+    mp_obj_t fn;
+} m_atexit_node_t;
+
 // atexit(callback): Callback is called when sys.exit is called.
 static mp_obj_t mp_sys_atexit(mp_obj_t obj) {
-    mp_obj_t old = MP_STATE_VM(sys_exitfunc);
-    MP_STATE_VM(sys_exitfunc) = obj;
-    return old;
+    m_atexit_node_t *node = m_malloc(sizeof(m_atexit_node_t));
+    if (MP_STATE_VM(sys_exitfunc) != mp_const_none) {
+        MP_STATE_VM(sys_exitfunc)->prev = node;
+    }
+    node->fn = obj;
+    node->prev = NULL;
+    node->next = MP_STATE_VM(sys_exitfunc);
+    MP_STATE_VM(sys_exitfunc) = node;
+    return obj;
 }
 static MP_DEFINE_CONST_FUN_OBJ_1(mp_sys_atexit_obj, mp_sys_atexit);
+
+void mp_sys_atexit_execute(void) {
+    // walk through the linked list and execute each function
+    // Beware, the sys.settrace callback should be disabled before running sys.atexit.
+    while (MP_STATE_VM(sys_exitfunc) != mp_const_none) {
+        if (mp_obj_is_callable(MP_STATE_VM(sys_exitfunc)->fn)) {
+            nlr_buf_t nlr;
+            if (nlr_push(&nlr) == 0) {
+                mp_call_function_0(MP_STATE_VM(sys_exitfunc)->fn);
+            } else {
+                // Uncaught exception: print it out.
+                mp_obj_print_exception(&mp_plat_print, (mp_obj_t)nlr.ret_val);
+            }
+        }
+        MP_STATE_VM(sys_exitfunc) = MP_STATE_VM(sys_exitfunc)->next;
+    }
+}
 #endif
 
 #if MICROPY_PY_SYS_SETTRACE
@@ -364,7 +393,7 @@ MP_REGISTER_ROOT_POINTER(mp_obj_base_t * cur_exception);
 
 #if MICROPY_PY_SYS_ATEXIT
 // exposed through sys.atexit function
-MP_REGISTER_ROOT_POINTER(mp_obj_t sys_exitfunc);
+MP_REGISTER_ROOT_POINTER(struct _m_atexit_node_t *sys_exitfunc);
 #endif
 
 #if MICROPY_PY_SYS_ATTR_DELEGATION
diff --git a/py/mpconfig.h b/py/mpconfig.h
@@ -1488,7 +1488,7 @@ typedef double mp_float_t;
 
 // Whether to provide "sys.atexit" function (MicroPython extension)
 #ifndef MICROPY_PY_SYS_ATEXIT
-#define MICROPY_PY_SYS_ATEXIT (0)
+#define MICROPY_PY_SYS_ATEXIT (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EXTRA_FEATURES)
 #endif
 
 // Whether to provide the "sys.path" attribute (which forces module delegation
diff --git a/py/runtime.h b/py/runtime.h
@@ -95,6 +95,10 @@ extern const byte mp_binary_op_method_name[];
 void mp_init(void);
 void mp_deinit(void);
 
+#if MICROPY_PY_SYS_ATEXIT
+void mp_sys_atexit_execute(void);
+#endif
+
 void mp_sched_exception(mp_obj_t exc);
 void mp_sched_keyboard_interrupt(void);
 #if MICROPY_ENABLE_VM_ABORT
diff --git a/tests/misc/sys_atexit.py b/tests/misc/sys_atexit.py
@@ -15,6 +15,11 @@ def do_at_exit():
     print("done at exit:", some_var)
 
 
+@sys.atexit
+def do_at_exit_2():
+    print("done at exit last")
+
+
 sys.atexit(do_at_exit)
 
 some_var = "ok"
diff --git a/tests/misc/sys_atexit.py.exp b/tests/misc/sys_atexit.py.exp
@@ -1,2 +1,3 @@
 done before exit
 done at exit: ok
+done at exit last

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,3 @@`
`1`	`1`	`done before exit`
`2`	`2`	`done at exit: ok`
	`3`	`+done at exit last`