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 3434517328..3a91393812 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 a111b27aec..bded9da9fc 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 b7eb6bdeb1..1057cef041 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 70d7957198..04f9442b7b 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