esp8266/modmachine: Implement optional time_ms arg to machine.deepsleep.

1 files changed, 14 insertions, 2 deletions

diff --git a/ports/esp8266/modmachine.c b/ports/esp8266/modmachine.c
index 1c1d902e5..ba49fda88 100644
--- a/ports/esp8266/modmachine.c
+++ b/ports/esp8266/modmachine.c
@@ -102,7 +102,7 @@ STATIC mp_obj_t machine_sleep(void) {
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_0(machine_sleep_obj, machine_sleep);
 
-STATIC mp_obj_t machine_deepsleep(void) {
+STATIC mp_obj_t machine_deepsleep(size_t n_args, const mp_obj_t *args) {
     // default to sleep forever
     uint32_t sleep_us = 0;
 
@@ -123,6 +123,18 @@ STATIC mp_obj_t machine_deepsleep(void) {
         }
     }
 
+    // if an argument is given then that's the maximum time to sleep for
+    if (n_args == 1) {
+        mp_int_t max_ms = mp_obj_get_int(args[0]);
+        if (max_ms <= 0) {
+            max_ms = 1;
+        }
+        uint32_t max_us = max_ms * 1000;
+        if (sleep_us == 0 || max_us < sleep_us) {
+            sleep_us = max_us;
+        }
+    }
+
     // prepare for RTC reset at wake up
     rtc_prepare_deepsleep(sleep_us);
     // put the device in a deep-sleep state
@@ -136,7 +148,7 @@ STATIC mp_obj_t machine_deepsleep(void) {
 
     return mp_const_none;
 }
-STATIC MP_DEFINE_CONST_FUN_OBJ_0(machine_deepsleep_obj, machine_deepsleep);
+STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(machine_deepsleep_obj, 0, 1, machine_deepsleep);
 
 typedef struct _esp_timer_obj_t {
     mp_obj_base_t base;