Add supervisor.set_next_code_file()

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 ""
@@ -2462,7 +2466,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 ""
 

main.c

@@ -280,7 +280,12 @@ STATIC bool run_code_py(safe_mode_t safe_mode) {
     result.exception_type = NULL;
     result.exception_line = 0;
 
+    bool skip_repl;
+    bool skip_wait = false;
     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) {
         static const char * const supported_filenames[] = STRING_LIST(
@@ -302,25 +307,71 @@ STATIC bool run_code_py(safe_mode_t safe_mode) {
         usb_setup_with_vm();
         #endif
 
-        // This is where the user's python code is actually executed:
-        found_main = maybe_run_list(supported_filenames, &result);
-        // If that didn't work, double check the extensions
-        #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"));
+        // Check if a different run file has been allocated
+        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, ""};
+                // This is where the user's python code is actually executed:
+                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"));
+                }
             }
         }
-        #else
-        (void) found_main;
-        #endif
+        // Otherwise, default to the standard list of filenames
+        if (!found_main) {
+            // This is where the user's python code is actually executed:
+            found_main = maybe_run_list(supported_filenames, &result);
+            // If that didn't work, double check the extensions
+            #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"));
+                }
+            }
+            #else
+            (void) found_main;
+            #endif
+        }
 
         // Finished executing python code. Cleanup includes a board reset.
         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) {
+            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;
+                skip_wait = true;
+            }
+        }
+        else {
+            next_code_stickiness_situation |= SUPERVISOR_NEXT_CODE_OPT_STICKY_ON_ERROR;
+            // Deep sleep cannot be skipped
+            // TODO: settings in deep sleep should persist, using a new sleep memory API
+            if (next_code_options & SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_ERROR
+                && !(result.return_code & PYEXEC_DEEP_SLEEP)) {
+                skip_repl = true;
+                skip_wait = true;
+            }
+        }
         if (result.return_code & PYEXEC_FORCED_EXIT) {
-            return reload_requested;
+            skip_repl = reload_requested;
+            skip_wait = true;
         }
 
         if (reload_requested && result.return_code == PYEXEC_EXCEPTION) {
@@ -368,12 +419,17 @@ STATIC bool run_code_py(safe_mode_t safe_mode) {
     #if CIRCUITPY_ALARM
     bool fake_sleeping = false;
     #endif
-    bool skip_repl = false;
-    while (true) {
+    while (!skip_wait) {
         RUN_BACKGROUND_TASKS;
 
         // If a reload was requested by the supervisor or autoreload, return
         if (reload_requested) {
+            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).
             reload_requested = false;
             skip_repl = true;
             break;
@@ -508,6 +564,13 @@ STATIC bool run_code_py(safe_mode_t safe_mode) {
             port_idle_until_interrupt();
         }
     }
+
+    // free code allocation if unused
+    if ((next_code_options & next_code_stickiness_situation) == 0) {
+        free_memory(next_code_allocation);
+        next_code_allocation = NULL;
+    }
+
     // Done waiting, start the board back up.
     #if CIRCUITPY_STATUS_LED
     if (led_active) {
@@ -521,6 +584,7 @@ STATIC bool run_code_py(safe_mode_t safe_mode) {
         board_init();
     }
     #endif
+
     return skip_repl;
 }
 

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,98 @@ 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 +216,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);

supervisor/shared/autoreload.c

@@ -30,6 +30,7 @@
 #include "py/reload.h"
 #include "supervisor/shared/tick.h"
 
+supervisor_allocation *next_code_allocation;
 #include "shared-bindings/supervisor/Runtime.h"
 
 static volatile uint32_t autoreload_delay_ms = 0;

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);

supervisor/shared/memory.c

@@ -36,6 +36,8 @@ enum {
     CIRCUITPY_SUPERVISOR_IMMOVABLE_ALLOC_COUNT =
         // stack + heap
         2
+        // next_code_allocation
+        + 1
 
         #if INTERNAL_FLASH_FILESYSTEM == 0
         + 1