mpy: implement hacky disk mode

components/micropython/vendor/ports/esp32/main.c

@@ -78,6 +78,25 @@ int vprintf_null(const char *format, va_list ap) {
     return 0;
 }
 
+static void _diskmode_maybe(int ret) {
+    // Switch to disk mode and suspend forever. For now we leave disk
+    // mode by restarting the badge.
+    //
+    // TODO(q3k): handle disk mode exiting gracefully
+    if (ret & PYEXEC_DISK_MODE_SD) {
+        st3m_mode_set(st3m_mode_kind_disk_sd, NULL);
+        for (;;) {
+            vTaskDelay(100 / portTICK_PERIOD_MS);
+        }
+    }
+    if (ret & PYEXEC_DISK_MODE_FLASH) {
+        st3m_mode_set(st3m_mode_kind_disk_flash, NULL);
+        for (;;) {
+            vTaskDelay(100 / portTICK_PERIOD_MS);
+        }
+    }
+}
+
 void mp_task(void *pvParameter) {
     volatile uint32_t sp = (uint32_t)esp_cpu_get_sp();
     #if MICROPY_PY_THREAD
@@ -120,8 +139,10 @@ soft_reset:
     // run boot-up scripts
     pyexec_frozen_module("_boot.py");
     pyexec_file_if_exists("boot.py");
+    int ret;
     if (pyexec_mode_kind == PYEXEC_MODE_FRIENDLY_REPL) {
-        int ret = pyexec_file_if_exists("/flash/sys/main.py");
+        ret = pyexec_file_if_exists("/flash/sys/main.py");
+        _diskmode_maybe(ret);
         if (ret & PYEXEC_FORCED_EXIT) {
             goto soft_reset_exit;
         }
@@ -137,7 +158,8 @@ soft_reset:
             }
             esp_log_set_vprintf(vprintf_log);
         } else {
-            if (pyexec_friendly_repl() != 0) {
+            if ((ret = pyexec_friendly_repl() != 0)) {
+                _diskmode_maybe(pyexec_system_exit);
                 break;
             }
         }

components/micropython/vendor/ports/esp32/modmachine.c

@@ -233,6 +233,18 @@ STATIC mp_obj_t machine_soft_reset(void) {
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_0(machine_soft_reset_obj, machine_soft_reset);
 
+STATIC mp_obj_t machine_disk_mode_sd(void) {
+    pyexec_system_exit = PYEXEC_FORCED_EXIT | PYEXEC_DISK_MODE_SD;
+    mp_raise_type(&mp_type_SystemExit);
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_0(machine_disk_mode_sd_obj, machine_disk_mode_sd);
+
+STATIC mp_obj_t machine_disk_mode_flash(void) {
+    pyexec_system_exit = PYEXEC_FORCED_EXIT | PYEXEC_DISK_MODE_FLASH;
+    mp_raise_type(&mp_type_SystemExit);
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_0(machine_disk_mode_flash_obj, machine_disk_mode_flash);
+
 STATIC mp_obj_t machine_unique_id(void) {
     uint8_t chipid[6];
     esp_efuse_mac_get_default(chipid);
@@ -271,6 +283,8 @@ STATIC const mp_rom_map_elem_t machine_module_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR_freq), MP_ROM_PTR(&machine_freq_obj) },
     { MP_ROM_QSTR(MP_QSTR_reset), MP_ROM_PTR(&machine_reset_obj) },
     { MP_ROM_QSTR(MP_QSTR_soft_reset), MP_ROM_PTR(&machine_soft_reset_obj) },
+    { MP_ROM_QSTR(MP_QSTR_disk_mode_sd), MP_ROM_PTR(&machine_disk_mode_sd_obj) },
+    { MP_ROM_QSTR(MP_QSTR_disk_mode_flash), MP_ROM_PTR(&machine_disk_mode_flash_obj) },
     { MP_ROM_QSTR(MP_QSTR_unique_id), MP_ROM_PTR(&machine_unique_id_obj) },
     { MP_ROM_QSTR(MP_QSTR_sleep), MP_ROM_PTR(&machine_lightsleep_obj) },
     { MP_ROM_QSTR(MP_QSTR_lightsleep), MP_ROM_PTR(&machine_lightsleep_obj) },

components/micropython/vendor/ports/esp32/modmachine.h

@@ -32,4 +32,7 @@ void machine_pwm_deinit_all(void);
 void machine_timer_deinit_all(void);
 void machine_i2s_init0();
 
+#define PYEXEC_DISK_MODE_SD (0x200)
+#define PYEXEC_DISK_MODE_FLASH (0x400)
+
 #endif // MICROPY_INCLUDED_ESP32_MODMACHINE_H

python_payload/mypystubs/machine.pyi

@@ -3,6 +3,18 @@ from typing import Any
 def reset() -> None:
     pass
 
+def disk_mode_sd() -> None:
+    """
+    Exit Micropython and enable SD card disk mode.
+    """
+    pass
+
+def disk_mode_flash() -> None:
+    """
+    Exit Micropython and enable internal flash disk mode.
+    """
+    pass
+
 class Pin:
     """
     See the official micropython docs for more details: