Add lots of options.

Author: cwalther

main.c

@@ -281,14 +281,8 @@ bool cleanup_after_vm(supervisor_allocation* heap) {
     return false;
 }
 
-void maybe_clear_next_code_info(bool new_next_code) {
-    // - If a new next code file was set, honor that in any case.
-    // - Otherwise: If a reload was requested by a USB file write, we want to run the same file
-    //   again, preserve any next-code info. Currently that also covers a reload requested by
-    //   supervisor.reload(), or do we want to treat that differently by turning
-    //   reload_requested into an enum?
-    // - Otherwise: Clear any leftover next-code info and restart at code.py/main.py.
-    if (!new_next_code && !reload_requested) {
+void maybe_clear_next_code_info(uint8_t reasons_to_keep) {
+    if (reasons_to_keep == 0) {
         free_memory(next_code_allocation);
         next_code_allocation = NULL;
     }
@@ -317,7 +311,9 @@ bool run_code_py(safe_mode_t safe_mode) {
     result.exception_line = 0;
 
     bool found_main = false;
-    bool new_next_code = false;
+    uint8_t next_code_options = 0;
+    // Collects stickiness bits that apply in the current situation.
+    uint8_t next_code_stickiness_situation = SUPERVISOR_NEXT_CODE_OPT_NEWLY_SET;
 
     if (safe_mode != NO_SAFE_MODE) {
         serial_write_compressed(translate("Running in safe mode! "));
@@ -335,7 +331,6 @@ bool run_code_py(safe_mode_t safe_mode) {
         filesystem_flush();
         supervisor_allocation* heap = allocate_remaining_memory();
         start_mp(heap);
-        uint8_t next_code_options = 0;
         if (next_code_allocation) {
             next_code_options = ((next_code_info_t*)next_code_allocation->ptr)->options;
             const char* next_list[] = {((next_code_info_t*)next_code_allocation->ptr)->filename, ""};
@@ -356,14 +351,31 @@ bool run_code_py(safe_mode_t safe_mode) {
             }
             #endif
         }
-        new_next_code = cleanup_after_vm(heap);
+        bool new_next_code = cleanup_after_vm(heap);
 
-        if (result.return_code == 0 && (next_code_options & SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_SUCCESS)) {
-            maybe_clear_next_code_info(new_next_code);
-            return true;
+        // If a new next code file was set, that is a reason to keep it (obviously). Stuff this into
+        // the options because it can be treated like any other reason-for-stickiness bit.
+        if (new_next_code) next_code_options |= SUPERVISOR_NEXT_CODE_OPT_NEWLY_SET;
+
+        if (reload_requested) {
+            next_code_stickiness_situation |= SUPERVISOR_NEXT_CODE_OPT_STICKY_ON_RELOAD;
+        }
+        if (result.return_code == 0) {
+            next_code_stickiness_situation |= SUPERVISOR_NEXT_CODE_OPT_STICKY_ON_SUCCESS;
+            if (next_code_options & SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_SUCCESS) {
+                maybe_clear_next_code_info(next_code_options & next_code_stickiness_situation);
+                return true;
+            }
+        }
+        else {
+            next_code_stickiness_situation |= SUPERVISOR_NEXT_CODE_OPT_STICKY_ON_ERROR;
+            if (next_code_options & SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_ERROR) {
+                maybe_clear_next_code_info(next_code_options & next_code_stickiness_situation);
+                return true;
+            }
         }
         if (result.return_code & PYEXEC_FORCED_EXIT) {
-            maybe_clear_next_code_info(new_next_code);
+            maybe_clear_next_code_info(next_code_options & next_code_stickiness_situation);
             return reload_requested;
         }
     }
@@ -382,13 +394,14 @@ bool run_code_py(safe_mode_t safe_mode) {
     while (true) {
         RUN_BACKGROUND_TASKS;
         if (reload_requested) {
-            // no need to call maybe_clear_next_code_info(), it would be a no-op
+            next_code_stickiness_situation |= SUPERVISOR_NEXT_CODE_OPT_STICKY_ON_RELOAD;
+            maybe_clear_next_code_info(next_code_options & next_code_stickiness_situation);
             reload_requested = false;
             return true;
         }
 
         if (serial_connected() && serial_bytes_available()) {
-            maybe_clear_next_code_info(new_next_code);
+            maybe_clear_next_code_info(next_code_options & next_code_stickiness_situation);
             // Skip REPL if reload was requested.
             return (serial_read() == CHAR_CTRL_D);
         }

shared-bindings/supervisor/__init__.c

@@ -110,23 +110,70 @@ STATIC mp_obj_t supervisor_set_next_stack_limit(mp_obj_t size_obj) {
 }
 MP_DEFINE_CONST_FUN_OBJ_1(supervisor_set_next_stack_limit_obj, supervisor_set_next_stack_limit);
 
+//| def set_next_code_file(filename: Optional[str], /, *, reload_on_success : bool = False, reload_on_error: bool = False, sticky_on_success: bool = False, sticky_on_error: bool = False, sticky_on_reload: bool = False) -> None:
+//|     """Set what file to run on the next vm run.
+//|
+//|     When not None, the given filename is inserted at the front of the usual ['code.py',
+//|     'main.py'] search sequence.
+//|
+//|     The optional keyword arguments specify what happens after the specified file has run:
+//|
+//|     sticky_on_… determine whether the newly set filename and options stay in effect: If True,
+//|     further runs will continue to run that file (unless it says otherwise by calling
+//|     set_next_code_filename() itself). If False, the settings will only affect one run and revert
+//|     to the standard code.py/main.py afterwards.
+//|
+//|     reload_on_… determine how to continue: If False, wait in the usual "Code done running.
+//|     Waiting for reload. / Press any key to enter the REPL. Use CTRL-D to reload." state. If
+//|     True, reload immediately as if CTRL-D was pressed.
+//|
+//|     …_on_success take effect when the program runs to completion or calls sys.exit().
+//|
+//|     …_on_error take effect when the program exits with an exception, including the
+//|     KeyboardInterrupt caused by CTRL-C.
+//|
+//|     …_on_reload take effect when the program is interrupted by files being written to the USB
+//|     drive (auto-reload) or when it calls supervisor.reload().
+//|
+//|     These settings are stored in RAM, not in persistent memory, and will therefore only affect
+//|     soft reloads. Powering off or resetting the device will always revert to standard settings.
+//|
+//|     When called multiple times in the same run, only the last call takes effect, replacing any
+//|     settings made by previous ones. This is the main use of passing None as a filename: to
+//|     reset to the standard search sequence."""
+//|     ...
+//|
 STATIC mp_obj_t supervisor_set_next_code_file(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
     static const mp_arg_t allowed_args[] = {
         { MP_QSTR_filename, MP_ARG_REQUIRED | MP_ARG_OBJ, {.u_rom_obj = mp_const_none} },
         { MP_QSTR_reload_on_success, MP_ARG_KW_ONLY | MP_ARG_BOOL, {.u_bool = false} },
+        { MP_QSTR_reload_on_error, MP_ARG_KW_ONLY | MP_ARG_BOOL, {.u_bool = false} },
+        { MP_QSTR_sticky_on_success, MP_ARG_KW_ONLY | MP_ARG_BOOL, {.u_bool = false} },
+        { MP_QSTR_sticky_on_error, MP_ARG_KW_ONLY | MP_ARG_BOOL, {.u_bool = false} },
+        { MP_QSTR_sticky_on_reload, MP_ARG_KW_ONLY | MP_ARG_BOOL, {.u_bool = false} },
     };
     struct {
         mp_arg_val_t filename;
         mp_arg_val_t reload_on_success;
+        mp_arg_val_t reload_on_error;
+        mp_arg_val_t sticky_on_success;
+        mp_arg_val_t sticky_on_error;
+        mp_arg_val_t sticky_on_reload;
     } args;
     mp_arg_parse_all(n_args, pos_args, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, (mp_arg_val_t*)&args);
     if (!MP_OBJ_IS_STR_OR_BYTES(args.filename.u_obj) && args.filename.u_obj != mp_const_none) {
         mp_raise_TypeError(translate("argument has wrong type"));
     }
     if (args.filename.u_obj == mp_const_none) args.filename.u_obj = mp_const_empty_bytes;
+    uint8_t options = 0;
+    if (args.reload_on_success.u_bool) options |= SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_SUCCESS;
+    if (args.reload_on_error.u_bool) options |= SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_ERROR;
+    if (args.sticky_on_success.u_bool) options |= SUPERVISOR_NEXT_CODE_OPT_STICKY_ON_SUCCESS;
+    if (args.sticky_on_error.u_bool) options |= SUPERVISOR_NEXT_CODE_OPT_STICKY_ON_ERROR;
+    if (args.sticky_on_reload.u_bool) options |= SUPERVISOR_NEXT_CODE_OPT_STICKY_ON_RELOAD;
     mp_obj_t items[] = {
         args.filename.u_obj,
-        args.reload_on_success.u_bool ? MP_OBJ_NEW_SMALL_INT(SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_SUCCESS) : MP_OBJ_NEW_SMALL_INT(0)
+        MP_OBJ_NEW_SMALL_INT(options)
     };
     MP_STATE_VM(supervisor_next_code) = mp_obj_new_tuple(MP_ARRAY_SIZE(items), items);
     return mp_const_none;

supervisor/shared/autoreload.h

@@ -30,7 +30,12 @@
 #include <stdbool.h>
 
 enum {
-	SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_SUCCESS = 0x1
+	SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_SUCCESS = 0x1,
+	SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_ERROR = 0x2,
+	SUPERVISOR_NEXT_CODE_OPT_STICKY_ON_SUCCESS = 0x4,
+	SUPERVISOR_NEXT_CODE_OPT_STICKY_ON_ERROR = 0x8,
+	SUPERVISOR_NEXT_CODE_OPT_STICKY_ON_RELOAD = 0x10,
+	SUPERVISOR_NEXT_CODE_OPT_NEWLY_SET = 0x20,
 };
 
 extern volatile bool reload_requested;