atexit module refinements

microdev1 · microdev1 · commit 1c4a6c3667ec · 2021-08-16T21:37:32.000+05:30
- add test for atexit module
- add callback to gc collection
- fix callback memory allocation
- execute callback on both code and repl exit
diff --git a/lib/utils/pyexec.c b/lib/utils/pyexec.c
@@ -44,10 +44,6 @@
 #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;
 
@@ -131,10 +127,6 @@ STATIC int parse_compile_execute(const void *source, mp_parse_input_kind_t input
         start = mp_hal_ticks_ms();
         #endif
         mp_call_function_0(module_fun);
-        // execute exit handlers (if any).
-        #if CIRCUITPY_ATEXIT
-        shared_module_atexit_execute();
-        #endif
         mp_hal_set_interrupt_char(-1); // disable interrupt
         mp_handle_pending(true); // handle any pending exceptions (and any callbacks)
         nlr_pop();
diff --git a/main.c b/main.c
@@ -217,6 +217,9 @@ STATIC bool maybe_run_list(const char * const * filenames, pyexec_result_t* exec
     decompress(compressed, decompressed);
     mp_hal_stdout_tx_str(decompressed);
     pyexec_file(filename, exec_result);
+    #if CIRCUITPY_ATEXIT
+    shared_module_atexit_execute();
+    #endif
     return true;
 }
 
@@ -766,6 +769,9 @@ STATIC int run_repl(void) {
     } else {
         exit_code = pyexec_friendly_repl();
     }
+    #if CIRCUITPY_ATEXIT
+    shared_module_atexit_execute();
+    #endif
     cleanup_after_vm(heap, MP_OBJ_SENTINEL);
     #if CIRCUITPY_STATUS_LED
     status_led_init();
@@ -881,6 +887,10 @@ void gc_collect(void) {
     common_hal_alarm_gc_collect();
     #endif
 
+    #if CIRCUITPY_ATEXIT
+    atexit_gc_collect();
+    #endif
+
     #if CIRCUITPY_DISPLAYIO
     displayio_gc_collect();
     #endif
diff --git a/ports/atmel-samd/boards/adafruit_proxlight_trinkey_m0/mpconfigboard.mk b/ports/atmel-samd/boards/adafruit_proxlight_trinkey_m0/mpconfigboard.mk
@@ -20,7 +20,6 @@ CIRCUITPY_ROTARYIO = 0
 CIRCUITPY_RTC = 0
 CIRCUITPY_USB_MIDI = 0
 
-CIRCUITPY_ATEXIT = 0
 CIRCUITPY_PIXELBUF = 1
 CIRCUITPY_BUSDEVICE = 1
 
diff --git a/py/circuitpy_mpconfig.mk b/py/circuitpy_mpconfig.mk
@@ -48,7 +48,7 @@ CFLAGS += -DCIRCUITPY_ALARM=$(CIRCUITPY_ALARM)
 CIRCUITPY_ANALOGIO ?= 1
 CFLAGS += -DCIRCUITPY_ANALOGIO=$(CIRCUITPY_ANALOGIO)
 
-CIRCUITPY_ATEXIT ?= 1
+CIRCUITPY_ATEXIT ?= $(CIRCUITPY_FULL_BUILD)
 CFLAGS += -DCIRCUITPY_ATEXIT=$(CIRCUITPY_ATEXIT)
 
 CIRCUITPY_AUDIOBUSIO ?= $(CIRCUITPY_FULL_BUILD)
diff --git a/shared-bindings/atexit/__init__.c b/shared-bindings/atexit/__init__.c
@@ -31,6 +31,9 @@
 //| This module defines functions to register and unregister cleanup functions.
 //| Functions thus registered are automatically executed upon normal vm termination.
 //|
+//| These functions are run in the reverse order in which they were registered;
+//| if you register ``A``, ``B``, and ``C``, they will be run in the order ``C``, ``B``, ``A``.
+//|
 //| """
 //| ...
 //|
@@ -39,13 +42,14 @@
 //|
 //|     """Register func as a function to be executed at termination.
 //|
-//|        Any optional arguments that are to be passed to func must be passed as arguments to register().
+//|        Any optional arguments that are to be passed to func must be passed as arguments to `register()`.
 //|        It is possible to register the same function and arguments more than once.
 //|
 //|        At normal program termination (for instance, if `sys.exit()` is called or the vm execution completes),
-//|        all functions registered are called in order of there registration.
+//|        all functions registered are called in last in, first out order.
 //|
-//|        If an exception is raised during execution of the exit handlers, a traceback is printed and the execution stops.
+//|        If an exception is raised during execution of the exit handler,
+//|        a traceback is printed (unless `SystemExit` is raised) and the execution stops.
 //|
 //|        This function returns func, which makes it possible to use it as a decorator.
 //|
diff --git a/shared-module/atexit/__init__.c b/shared-module/atexit/__init__.c
@@ -24,6 +24,7 @@
  * THE SOFTWARE.
  */
 
+#include "py/gc.h"
 #include "py/runtime.h"
 #include "shared-module/atexit/__init__.h"
 
@@ -32,15 +33,19 @@ typedef struct _atexit_callback_t {
     mp_obj_t func, *args;
 } atexit_callback_t;
 
-size_t callback_len = 0;
-atexit_callback_t *callback = NULL;
+static size_t callback_len = 0;
+static atexit_callback_t *callback = NULL;
 
 void atexit_reset(void) {
     callback_len = 0;
     m_free(callback);
     callback = NULL;
 }
 
+void atexit_gc_collect(void) {
+    gc_collect_ptr(callback);
+}
+
 void shared_module_atexit_register(mp_obj_t *func, size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
     if (!mp_obj_is_callable(func)) {
         mp_raise_TypeError_varg(translate("'%q' object is not callable"), mp_obj_get_type_qstr(func));
@@ -50,24 +55,22 @@ void shared_module_atexit_register(mp_obj_t *func, size_t n_args, const mp_obj_t
         .n_pos = 0,
         .n_kw = 0,
         .func = func,
-        .args = ((n_args + n_kw_args) == 0) ? NULL : m_malloc((n_args + (2 * n_kw_args)) * sizeof(mp_obj_t), false)
+        .args = (n_args + n_kw_args) ? m_malloc((n_args + (n_kw_args * 2)) * sizeof(mp_obj_t), false) : NULL
     };
     for (; cb.n_pos < n_args; cb.n_pos++) {
         cb.args[cb.n_pos] = pos_args[cb.n_pos];
     }
     for (size_t i = cb.n_pos; cb.n_kw < n_kw_args; i++, cb.n_kw++) {
-        cb.args[i] = kw_args[cb.n_kw].table->key;
-        cb.args[i += 1] = kw_args[cb.n_kw].table->value;
-    }
-    if (!callback) {
-        callback = (atexit_callback_t *)m_realloc(callback, sizeof(cb));
+        cb.args[i] = kw_args->table[cb.n_kw].key;
+        cb.args[i += 1] = kw_args->table[cb.n_kw].value;
     }
+    callback = (atexit_callback_t *)m_realloc(callback, (callback_len + 1) * sizeof(cb));
     callback[callback_len++] = cb;
 }
 
 void shared_module_atexit_unregister(const mp_obj_t *func) {
     for (size_t i = 0; i < callback_len; i++) {
-        if (callback[i].func == func) {
+        if (callback[i].func == *func) {
             callback[i].n_pos = 0;
             callback[i].n_kw = 0;
             callback[i].func = mp_const_none;
@@ -78,7 +81,7 @@ void shared_module_atexit_unregister(const mp_obj_t *func) {
 
 void shared_module_atexit_execute(void) {
     if (callback) {
-        for (size_t i = 0; i < callback_len; i++) {
+        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);
             }
diff --git a/shared-module/atexit/__init__.h b/shared-module/atexit/__init__.h
@@ -30,6 +30,7 @@
 #include "py/obj.h"
 
 extern void atexit_reset(void);
+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);
diff --git a/tests/circuitpython/atexit_test.py b/tests/circuitpython/atexit_test.py
@@ -0,0 +1,22 @@
+# test atexit module
+
+try:
+    import atexit
+except ImportError:
+    print("SKIP")
+    raise SystemExit
+
+
+@atexit.register
+def skip_at_exit():
+    print("skip at exit")
+
+
+@atexit.register
+def do_at_exit(*args, **kwargs):
+    print("done at exit:", args, kwargs)
+
+
+atexit.unregister(skip_at_exit)
+atexit.register(do_at_exit, "ok", 1, arg="2", param=(3, 4))
+print("done before exit")
