diff --git a/ports/unix/main.c b/ports/unix/main.c index b69df90354..55bdc22082 100644 --- a/ports/unix/main.c +++ b/ports/unix/main.c @@ -690,6 +690,13 @@ MP_NOINLINE int main_(int argc, char **argv) { MP_STATE_THREAD(prof_trace_callback) = MP_OBJ_NULL; #endif + #if MICROPY_PY_SYS_ATEXIT + // Beware, the sys.settrace callback should be disabled before running sys.atexit. + if (mp_obj_is_callable(MP_STATE_VM(sys_exitfunc))) { + mp_call_function_0(MP_STATE_VM(sys_exitfunc)); + } + #endif + #if MICROPY_PY_MICROPYTHON_MEM_INFO if (mp_verbose_flag) { mp_micropython_mem_info(0, NULL); diff --git a/py/modsys.c b/py/modsys.c index 1e31a04d27..cbd1712920 100644 --- a/py/modsys.c +++ b/py/modsys.c @@ -166,6 +166,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 + #if MICROPY_PY_SYS_SETTRACE // settrace(tracefunc): Set the system’s trace function. STATIC mp_obj_t mp_sys_settrace(mp_obj_t obj) { @@ -227,6 +237,14 @@ STATIC const mp_rom_map_elem_t mp_module_sys_globals_table[] = { #if MICROPY_PY_SYS_GETSIZEOF { MP_ROM_QSTR(MP_QSTR_getsizeof), MP_ROM_PTR(&mp_sys_getsizeof_obj) }, #endif + + /* + * Extensions to CPython + */ + + #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); diff --git a/py/mpconfig.h b/py/mpconfig.h index 8419a50d5e..2cd2a0d356 100644 --- a/py/mpconfig.h +++ b/py/mpconfig.h @@ -1311,6 +1311,11 @@ typedef double mp_float_t; #define MICROPY_PY_SYS_EXIT (1) #endif +// Whether to provide "sys.atexit" function (MicroPython extension) +#ifndef MICROPY_PY_SYS_ATEXIT +#define MICROPY_PY_SYS_ATEXIT (0) +#endif + // Whether to provide "sys.settrace" function #ifndef MICROPY_PY_SYS_SETTRACE #define MICROPY_PY_SYS_SETTRACE (0) diff --git a/py/mpstate.h b/py/mpstate.h index e868a773bf..423463109d 100644 --- a/py/mpstate.h +++ b/py/mpstate.h @@ -157,6 +157,11 @@ typedef struct _mp_state_vm_t { mp_obj_base_t *cur_exception; #endif + #if MICROPY_PY_SYS_ATEXIT + // exposed through sys.atexit function + mp_obj_t sys_exitfunc; + #endif + // dictionary for the __main__ module mp_obj_dict_t dict_main; diff --git a/py/runtime.c b/py/runtime.c index ba0d59dffa..ceb5e83b14 100644 --- a/py/runtime.c +++ b/py/runtime.c @@ -130,6 +130,10 @@ void mp_init(void) { sizeof(MP_STATE_VM(fs_user_mount)) - MICROPY_FATFS_NUM_PERSISTENT); #endif + #if MICROPY_PY_SYS_ATEXIT + MP_STATE_VM(sys_exitfunc) = mp_const_none; + #endif + #if MICROPY_PY_SYS_SETTRACE MP_STATE_THREAD(prof_trace_callback) = MP_OBJ_NULL; MP_STATE_THREAD(prof_callback_is_executing) = false;