py: Implement new sys.atexit feature.

This patch implements a new sys.atexit function which registers a function
that is later executed when the main script ends.  It is configurable via
MICROPY_PY_SYS_ATEXIT, disabled by default.

This is not compliant with CPython, rather it can be used to implement a
CPython compatible "atexit" module if desired (similar to how
sys.print_exception can be used to implement functionality of the
"traceback" module).

1 files changed, 13 insertions, 0 deletions

diff --git a/py/modsys.c b/py/modsys.c
index 343451732..3a9139381 100644
--- a/py/modsys.c
+++ b/py/modsys.c
@@ -146,6 +146,16 @@ STATIC mp_obj_t mp_sys_getsizeof(mp_obj_t obj) {
 STATIC MP_DEFINE_CONST_FUN_OBJ_1(mp_sys_getsizeof_obj, mp_sys_getsizeof);
 #endif
 
+#if MICROPY_PY_SYS_ATEXIT
+// 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;
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_1(mp_sys_atexit_obj, mp_sys_atexit);
+#endif
+
 STATIC const mp_rom_map_elem_t mp_module_sys_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_sys) },
 
@@ -201,6 +211,9 @@ STATIC const mp_rom_map_elem_t mp_module_sys_globals_table[] = {
      */
 
     { MP_ROM_QSTR(MP_QSTR_print_exception), MP_ROM_PTR(&mp_sys_print_exception_obj) },
+    #if MICROPY_PY_SYS_ATEXIT
+    { MP_ROM_QSTR(MP_QSTR_atexit), MP_ROM_PTR(&mp_sys_atexit_obj) },
+    #endif
 };
 
 STATIC MP_DEFINE_CONST_DICT(mp_module_sys_globals, mp_module_sys_globals_table);