adafruit · tannewt · Jun 25, 2021 · Oct 29, 2020 · Jun 20, 2021 · Jun 20, 2021
diff --git a/locale/circuitpython.pot b/locale/circuitpython.pot
@@ -47,6 +47,10 @@ msgstr ""
 msgid " is of type %q\n"
 msgstr ""
 
+#: main.c
+msgid " not found.\n"
+msgstr ""
+
 #: main.c
 msgid " output:\n"
 msgstr ""
@@ -2213,7 +2217,7 @@ msgstr ""
 msgid "argsort is not implemented for flattened arrays"
 msgstr ""
 
-#: py/runtime.c
+#: py/runtime.c shared-bindings/supervisor/__init__.c
 msgid "argument has wrong type"
 msgstr ""
 

diff --git a/main.c b/main.c
@@ -272,7 +272,11 @@ STATIC bool run_code_py(safe_mode_t safe_mode) {
     result.exception_type = NULL;
     result.exception_line = 0;
 
+    bool skip_repl;
     bool found_main = 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) {
         new_status_color(MAIN_RUNNING);
@@ -290,22 +294,62 @@ STATIC bool run_code_py(safe_mode_t safe_mode) {
         supervisor_allocation* heap = allocate_remaining_memory();
         start_mp(heap);
 
-        found_main = maybe_run_list(supported_filenames, &result);
-        #if CIRCUITPY_FULL_BUILD
-        if (!found_main){
-            found_main = maybe_run_list(double_extension_filenames, &result);
-            if (found_main) {
-                serial_write_compressed(translate("WARNING: Your code filename has two extensions\n"));
+        if (next_code_allocation) {
+            ((next_code_info_t*)next_code_allocation->ptr)->options &= ~SUPERVISOR_NEXT_CODE_OPT_NEWLY_SET;
+            next_code_options = ((next_code_info_t*)next_code_allocation->ptr)->options;
+            if (((next_code_info_t*)next_code_allocation->ptr)->filename[0] != '\0') {
+                const char* next_list[] = {((next_code_info_t*)next_code_allocation->ptr)->filename, ""};
+                found_main = maybe_run_list(next_list, &result);
+                if (!found_main) {
+                    serial_write(((next_code_info_t*)next_code_allocation->ptr)->filename);
+                    serial_write_compressed(translate(" not found.\n"));
+                }
             }
         }
-        #endif
+        if (!found_main) {
+            found_main = maybe_run_list(supported_filenames, &result);
+            #if CIRCUITPY_FULL_BUILD
+            if (!found_main){
+                found_main = maybe_run_list(double_extension_filenames, &result);
+                if (found_main) {
+                    serial_write_compressed(translate("WARNING: Your code filename has two extensions\n"));
+                }
+            }
+            #endif
+        }
 
         // TODO: on deep sleep, make sure display is refreshed before sleeping (for e-ink).
 
         cleanup_after_vm(heap);
 
+        // 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. The
+        // source is different though: it comes from the options that will apply to the next run,
+        // while the rest of next_code_options is what applied to this run.
+        if (next_code_allocation != NULL && (((next_code_info_t*)next_code_allocation->ptr)->options & SUPERVISOR_NEXT_CODE_OPT_NEWLY_SET)) {
+            next_code_options |= SUPERVISOR_NEXT_CODE_OPT_NEWLY_SET;
+        }
+
+        if (reload_requested) {
+            next_code_stickiness_situation |= SUPERVISOR_NEXT_CODE_OPT_STICKY_ON_RELOAD;
+        }
+        else if (result.return_code == 0) { //TODO mask out PYEXEC_DEEP_SLEEP?
+            next_code_stickiness_situation |= SUPERVISOR_NEXT_CODE_OPT_STICKY_ON_SUCCESS;
+            if (next_code_options & SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_SUCCESS) {
+                skip_repl = true;
+                goto done;
+            }
+        }
+        else {
+            next_code_stickiness_situation |= SUPERVISOR_NEXT_CODE_OPT_STICKY_ON_ERROR;
+            if (next_code_options & SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_ERROR) {
+                skip_repl = true;
+                goto done;
+            }
+        }
         if (result.return_code & PYEXEC_FORCED_EXIT) {
-            return reload_requested;
+            skip_repl = reload_requested;
+            goto done;
         }
 
         if (reload_requested && result.return_code == PYEXEC_EXCEPTION) {
@@ -333,9 +377,16 @@ STATIC bool run_code_py(safe_mode_t safe_mode) {
                 board_init();
             }
             #endif
+            next_code_stickiness_situation |= SUPERVISOR_NEXT_CODE_OPT_STICKY_ON_RELOAD;
+            // Should the STICKY_ON_SUCCESS and STICKY_ON_ERROR bits be cleared in
+            // next_code_stickiness_situation? I can see arguments either way, but I'm deciding
+            // "no" for now, mainly because it's a bit less code. At this point, we have both a
+            // success or error and a reload, so let's have both of the respective options take
+            // effect (in OR combination).
             supervisor_set_run_reason(RUN_REASON_AUTO_RELOAD);
             reload_requested = false;
-            return true;
+            skip_repl = true;
+            goto done;
         }
 
         if (serial_connected() && serial_bytes_available()) {
@@ -345,11 +396,11 @@ STATIC bool run_code_py(safe_mode_t safe_mode) {
             }
             #endif
             // Skip REPL if reload was requested.
-            bool ctrl_d = serial_read() == CHAR_CTRL_D;
-            if (ctrl_d) {
+            skip_repl = serial_read() == CHAR_CTRL_D;
+            if (skip_repl) {
                 supervisor_set_run_reason(RUN_REASON_REPL_RELOAD);
             }
-            return ctrl_d;
+            goto done;
         }
 
         // Check for a deep sleep alarm and restart the VM. This can happen if
@@ -428,6 +479,13 @@ STATIC bool run_code_py(safe_mode_t safe_mode) {
             port_idle_until_interrupt();
         }
     }
+
+done:
+    if ((next_code_options & next_code_stickiness_situation) == 0) {
+        free_memory(next_code_allocation);
+        next_code_allocation = NULL;
+    }
+    return skip_repl;
 }
 
 FIL* boot_output_file;

diff --git a/shared-bindings/supervisor/__init__.c b/shared-bindings/supervisor/__init__.c
@@ -23,9 +23,12 @@
  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  * THE SOFTWARE.
  */
+#include <string.h>
+
 #include "py/obj.h"
 #include "py/runtime.h"
 #include "py/reload.h"
+#include "py/objstr.h"
 
 #include "lib/utils/interrupt_char.h"
 #include "supervisor/shared/autoreload.h"
@@ -112,6 +115,87 @@ 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;
+    size_t len;
+    const char* filename = mp_obj_str_get_data(args.filename.u_obj, &len);
+    free_memory(next_code_allocation);
+    if (options != 0 || len != 0) {
+        next_code_allocation = allocate_memory(align32_size(sizeof(next_code_info_t) + len + 1), false, true);
+        if (next_code_allocation == NULL) {
+            m_malloc_fail(sizeof(next_code_info_t) + len + 1);
+        }
+        next_code_info_t* next_code = (next_code_info_t*)next_code_allocation->ptr;
+        next_code->options = options | SUPERVISOR_NEXT_CODE_OPT_NEWLY_SET;
+        memcpy(&next_code->filename, filename, len);
+        next_code->filename[len] = '\0';
+    }
+    else {
+        next_code_allocation = NULL;
+    }
+    return mp_const_none;
+}
+MP_DEFINE_CONST_FUN_OBJ_KW(supervisor_set_next_code_file_obj, 0, supervisor_set_next_code_file);
+
 STATIC const mp_rom_map_elem_t supervisor_module_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_supervisor) },
     { MP_ROM_QSTR(MP_QSTR_enable_autoreload),  MP_ROM_PTR(&supervisor_enable_autoreload_obj) },
@@ -121,7 +205,7 @@ STATIC const mp_rom_map_elem_t supervisor_module_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR_reload),  MP_ROM_PTR(&supervisor_reload_obj) },
     { MP_ROM_QSTR(MP_QSTR_RunReason),  MP_ROM_PTR(&supervisor_run_reason_type) },
     { MP_ROM_QSTR(MP_QSTR_set_next_stack_limit),  MP_ROM_PTR(&supervisor_set_next_stack_limit_obj) },
-
+    { MP_ROM_QSTR(MP_QSTR_set_next_code_file),  MP_ROM_PTR(&supervisor_set_next_code_file_obj) },
 };
 
 STATIC MP_DEFINE_CONST_DICT(supervisor_module_globals, supervisor_module_globals_table);

diff --git a/supervisor/shared/autoreload.c b/supervisor/shared/autoreload.c
@@ -30,6 +30,8 @@
 #include "py/reload.h"
 #include "supervisor/shared/tick.h"
 
+supervisor_allocation* next_code_allocation;
+
 static volatile uint32_t autoreload_delay_ms = 0;
 static bool autoreload_enabled = false;
 static bool autoreload_suspended = false;

diff --git a/supervisor/shared/autoreload.h b/supervisor/shared/autoreload.h
@@ -29,6 +29,24 @@
 
 #include <stdbool.h>
 
+#include "supervisor/memory.h"
+
+enum {
+	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,
+};
+
+typedef struct {
+    uint8_t options;
+    char filename[];
+} next_code_info_t;
+
+extern supervisor_allocation* next_code_allocation;
+
 extern volatile bool reload_requested;
 
 void autoreload_tick(void);

diff --git a/supervisor/shared/memory.c b/supervisor/shared/memory.c
@@ -36,6 +36,8 @@ enum {
     CIRCUITPY_SUPERVISOR_IMMOVABLE_ALLOC_COUNT =
     // stack + heap
     2
+    // next_code_allocation
+    + 1
 #ifdef EXTERNAL_FLASH_DEVICES
     + 1
 #endif