From cb3647004fbf7b069a3509a02b72420828369bd3 Mon Sep 17 00:00:00 2001
From: Milan Rossa <rossa.milan@gmail.com>
Date: Mon, 5 Aug 2019 15:06:41 +0200
Subject: [PATCH] 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).
---
py/modsys.c | 13 +++++++++++++
py/mpconfig.h | 5 +++++
py/mpstate.h | 5 +++++
py/runtime.c | 4 ++++
4 files changed, 27 insertions(+)
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);
diff --git a/py/mpconfig.h b/py/mpconfig.h
index a111b27ae..bded9da9f 100644
--- a/py/mpconfig.h
+++ b/py/mpconfig.h
@@ -1161,6 +1161,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.getsizeof" function
#ifndef MICROPY_PY_SYS_GETSIZEOF
#define MICROPY_PY_SYS_GETSIZEOF (0)
diff --git a/py/mpstate.h b/py/mpstate.h
index b7eb6bdeb..1057cef04 100644
--- a/py/mpstate.h
+++ b/py/mpstate.h
@@ -149,6 +149,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 70d795719..04f9442b7 100644
--- a/py/runtime.c
+++ b/py/runtime.c
@@ -118,6 +118,10 @@ void mp_init(void) {
MP_STATE_VM(vfs_mount_table) = NULL;
#endif
+ #if MICROPY_PY_SYS_ATEXIT
+ MP_STATE_VM(sys_exitfunc) = mp_const_none;
+ #endif
+
#if MICROPY_PY_THREAD_GIL
mp_thread_mutex_init(&MP_STATE_VM(gil_mutex));
#endif
--
GitLab