Merge pull request #5077 from microDev1/atexit

Add atexit module

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();
@@ -746,6 +764,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);

locale/circuitpython.pot

@@ -189,7 +189,7 @@ msgstr ""
 msgid "'%q' object is not an iterator"
 msgstr ""
 
-#: py/objtype.c py/runtime.c
+#: py/objtype.c py/runtime.c shared-module/atexit/__init__.c
 msgid "'%q' object is not callable"
 msgstr ""
 

main.c

@@ -70,6 +70,10 @@
 #include "shared-bindings/alarm/__init__.h"
 #endif
 
+#if CIRCUITPY_ATEXIT
+#include "shared-module/atexit/__init__.h"
+#endif
+
 #if CIRCUITPY_BLEIO
 #include "shared-bindings/_bleio/__init__.h"
 #include "supervisor/shared/bluetooth/bluetooth.h"
@@ -213,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(exec_result);
+    #endif
     return true;
 }
 
@@ -253,6 +260,11 @@ STATIC void cleanup_after_vm(supervisor_allocation* heap, mp_obj_t exception) {
 
     // Reset port-independent devices, like CIRCUITPY_BLEIO_HCI.
     reset_devices();
+
+    #if CIRCUITPY_ATEXIT
+    atexit_reset();
+    #endif
+
     // Turn off the display and flush the filesystem before the heap disappears.
     #if CIRCUITPY_DISPLAYIO
     reset_displays();
@@ -759,6 +771,13 @@ STATIC int run_repl(void) {
     } else {
         exit_code = pyexec_friendly_repl();
     }
+    #if CIRCUITPY_ATEXIT
+    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
     status_led_init();
@@ -875,6 +894,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

ports/unix/mpconfigport.h

@@ -107,7 +107,6 @@
 #define MICROPY_PY_BUILTINS_SLICE_ATTRS (1)
 #define MICROPY_PY_BUILTINS_SLICE_INDICES (1)
 #define MICROPY_PY_SYS_EXIT         (1)
-#define MICROPY_PY_SYS_ATEXIT       (1)
 #if MICROPY_PY_SYS_SETTRACE
 #define MICROPY_PERSISTENT_CODE_SAVE (1)
 #define MICROPY_COMP_CONST (0)

py/circuitpy_defns.mk

@@ -114,6 +114,9 @@ endif
 ifeq ($(CIRCUITPY_ANALOGIO),1)
 SRC_PATTERNS += analogio/%
 endif
+ifeq ($(CIRCUITPY_ATEXIT),1)
+SRC_PATTERNS += atexit/%
+endif
 ifeq ($(CIRCUITPY_AUDIOBUSIO),1)
 SRC_PATTERNS += audiobusio/%
 endif
@@ -485,6 +488,7 @@ SRC_SHARED_MODULE_ALL = \
 	_stage/__init__.c \
 	aesio/__init__.c \
 	aesio/aes.c \
+	atexit/__init__.c \
 	audiocore/RawSample.c \
 	audiocore/WaveFile.c \
 	audiocore/__init__.c \

py/circuitpy_mpconfig.h

@@ -265,6 +265,13 @@ extern const struct _mp_obj_module_t analogio_module;
 #define ANALOGIO_MODULE
 #endif
 
+#if CIRCUITPY_ATEXIT
+extern const struct _mp_obj_module_t atexit_module;
+#define ATEXIT_MODULE           { MP_OBJ_NEW_QSTR(MP_QSTR_atexit), (mp_obj_t)&atexit_module },
+#else
+#define ATEXIT_MODULE
+#endif
+
 #if CIRCUITPY_AUDIOBUSIO
 #define AUDIOBUSIO_MODULE        { MP_OBJ_NEW_QSTR(MP_QSTR_audiobusio), (mp_obj_t)&audiobusio_module },
 extern const struct _mp_obj_module_t audiobusio_module;
@@ -886,6 +893,7 @@ extern const struct _mp_obj_module_t msgpack_module;
     AESIO_MODULE \
     ALARM_MODULE \
     ANALOGIO_MODULE \
+    ATEXIT_MODULE \
     AUDIOBUSIO_MODULE \
     AUDIOCORE_MODULE \
     AUDIOIO_MODULE \

py/circuitpy_mpconfig.mk

@@ -48,6 +48,9 @@ CFLAGS += -DCIRCUITPY_ALARM=$(CIRCUITPY_ALARM)
 CIRCUITPY_ANALOGIO ?= 1
 CFLAGS += -DCIRCUITPY_ANALOGIO=$(CIRCUITPY_ANALOGIO)
 
+CIRCUITPY_ATEXIT ?= $(CIRCUITPY_FULL_BUILD)
+CFLAGS += -DCIRCUITPY_ATEXIT=$(CIRCUITPY_ATEXIT)
+
 CIRCUITPY_AUDIOBUSIO ?= $(CIRCUITPY_FULL_BUILD)
 CFLAGS += -DCIRCUITPY_AUDIOBUSIO=$(CIRCUITPY_AUDIOBUSIO)
 

shared-bindings/atexit/__init__.c

@@ -0,0 +1,93 @@
+/*
+ * This file is part of the Micro Python project, http://micropython.org/
+ *
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2021 microDev
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#include "shared-module/atexit/__init__.h"
+
+//| """Atexit Module
+//|
+//| 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``.
+//|
+//| """
+//| ...
+//|
+
+//| def register(func: Callable[..., Any], *args: Optional[Any], **kwargs: Optional[Any]) -> Callable[..., Any]:
+//|
+//|     """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()`.
+//|        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 last in, first out order.
+//|
+//|        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.
+//|
+//|     """
+//|     ...
+//|
+STATIC mp_obj_t atexit_register(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
+    shared_module_atexit_register(pos_args[0], (n_args - 1), ((n_args > 1) ? &pos_args[1] : NULL), kw_args);
+    return pos_args[0];
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_KW(atexit_register_obj, 1, atexit_register);
+
+//| def unregister(func: Callable[..., Any]) -> None:
+//|
+//|     """Remove func from the list of functions to be run at termination.
+//|
+//|        `unregister()` silently does nothing if func was not previously registered. If func has been registered more than once,
+//|        every occurrence of that function in the atexit call stack will be removed.
+//|
+//|     """
+//|     ...
+//|
+STATIC mp_obj_t atexit_unregister(const mp_obj_t self_in) {
+    shared_module_atexit_unregister(&self_in);
+    return mp_const_none;
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_1(atexit_unregister_obj, atexit_unregister);
+
+STATIC const mp_rom_map_elem_t atexit_module_globals_table[] = {
+    // module name
+    { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_atexit) },
+    // module functions
+    { MP_ROM_QSTR(MP_QSTR_register), MP_ROM_PTR(&atexit_register_obj) },
+    { MP_ROM_QSTR(MP_QSTR_unregister), MP_ROM_PTR(&atexit_unregister_obj) },
+};
+STATIC MP_DEFINE_CONST_DICT(atexit_module_globals, atexit_module_globals_table);
+
+const mp_obj_module_t atexit_module = {
+    .base = { &mp_type_module },
+    .globals = (mp_obj_dict_t *)&atexit_module_globals,
+};

shared-module/atexit/__init__.c

@@ -0,0 +1,92 @@
+/*
+ * This file is part of the Micro Python project, http://micropython.org/
+ *
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2021 microDev
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#include "py/gc.h"
+#include "py/runtime.h"
+#include "shared-module/atexit/__init__.h"
+
+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));
+    }
+    size_t n_kw_args = (kw_args) ? kw_args->used : 0;
+    atexit_callback_t cb = {
+        .n_pos = 0,
+        .n_kw = 0,
+        .func = func,
+        .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->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) {
+            callback[i].n_pos = 0;
+            callback[i].n_kw = 0;
+            callback[i].func = mp_const_none;
+            callback[i].args = NULL;
+        }
+    }
+}
+
+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) {
+                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

@@ -0,0 +1,46 @@
+/*
+ * This file is part of the Micro Python project, http://micropython.org/
+ *
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2021 microDev
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#ifndef MICROPY_INCLUDED_SHARED_MODULE_ATEXIT___INIT___H
+#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);
+
+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(pyexec_result_t *result);
+
+#endif  // MICROPY_INCLUDED_SHARED_MODULE_ATEXIT___INIT___H

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")

tests/misc/sys_atexit.py

@@ -1,21 +0,0 @@
-# test sys.atexit() function
-
-import sys
-
-try:
-    sys.atexit
-except AttributeError:
-    print("SKIP")
-    raise SystemExit
-
-some_var = None
-
-
-def do_at_exit():
-    print("done at exit:", some_var)
-
-
-sys.atexit(do_at_exit)
-
-some_var = "ok"
-print("done before exit")

tests/misc/sys_atexit.py.exp

@@ -1,2 +0,0 @@
-done before exit
-done at exit: ok

tests/unix/extra_coverage.py.exp

@@ -41,10 +41,10 @@ ime
 
 utime           utimeq
 
-argv            atexit          byteorder       exc_info
-exit            getsizeof       implementation  maxsize
-modules         path            platform        stderr
-stdin           stdout          version         version_info
+argv            byteorder       exc_info        exit
+getsizeof       implementation  maxsize         modules
+path            platform        stderr          stdin
+stdout          version         version_info
 ementation
 # attrtuple
 (start=1, stop=2, step=3)