pybricks.tools: Add hub_menu function.

laurensvalk · laurensvalk · commit ee5502f67473 · 2023-06-29T09:38:24.000+02:00
This implements the first approach discussed in pybricks/support#1064.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,7 @@
 - Added `'modes'` entry to the dictionary returned by `PUPDevice.info()`. It
   is a tuple of `(name, num_values, data_type)` tuples for each available mode.
 - Added `pybricks.tools.read_input_byte()` function ([support#1102]).
+- Added `pybricks.tools.hub_menu()` function ([support#1064]).
 
 ### Changed
 - Changed internal drivers for LEGO devices (motors and sensors) on all platforms.
@@ -16,6 +17,7 @@
 - Fixed hub will not power off when Bluetooth chip crashes on City and Technic hubs ([support#1095]).
 - Fixed `off()` method in `ColorLightMatrix`, `UltrasonicSensor`, `ColorSensor` ([support#1098]).
 
+[support#1064]: https://github.com/pybricks/support/issues/1064
 [support#1095]: https://github.com/pybricks/support/issues/1095
 [support#1098]: https://github.com/pybricks/support/issues/1098
 [support#1102]: https://github.com/pybricks/support/issues/1102
diff --git a/bricks/cityhub/mpconfigport.h b/bricks/cityhub/mpconfigport.h
@@ -41,6 +41,7 @@
 #define PYBRICKS_PY_ROBOTICS_DRIVEBASE_GYRO     (0)
 #define PYBRICKS_PY_ROBOTICS_DRIVEBASE_SPIKE    (0)
 #define PYBRICKS_PY_TOOLS                       (1)
+#define PYBRICKS_PY_TOOLS_HUB_MENU              (0)
 
 // Pybricks options
 #define PYBRICKS_OPT_COMPILER                   (1)
diff --git a/bricks/debug/mpconfigport.h b/bricks/debug/mpconfigport.h
@@ -36,6 +36,7 @@
 #define PYBRICKS_PY_PUPDEVICES                  (0)
 #define PYBRICKS_PY_ROBOTICS                    (0)
 #define PYBRICKS_PY_TOOLS                       (1)
+#define PYBRICKS_PY_TOOLS_HUB_MENU              (0)
 
 // Pybricks options
 #define PYBRICKS_OPT_COMPILER                   (0)
diff --git a/bricks/essentialhub/mpconfigport.h b/bricks/essentialhub/mpconfigport.h
@@ -42,6 +42,7 @@
 #define PYBRICKS_PY_ROBOTICS_DRIVEBASE_GYRO     (1)
 #define PYBRICKS_PY_ROBOTICS_DRIVEBASE_SPIKE    (1)
 #define PYBRICKS_PY_TOOLS                       (1)
+#define PYBRICKS_PY_TOOLS_HUB_MENU              (0)
 
 // Pybricks options
 #define PYBRICKS_OPT_COMPILER                   (1)
diff --git a/bricks/ev3dev/brickconfig.h b/bricks/ev3dev/brickconfig.h
@@ -44,4 +44,5 @@
 #define PYBRICKS_PY_ROBOTICS            (1)
 #define PYBRICKS_PY_ROBOTICS_DRIVEBASE_SPIKE (0)
 #define PYBRICKS_PY_TOOLS               (1)
+#define PYBRICKS_PY_TOOLS_HUB_MENU      (0)
 #define PYBRICKS_PY_USIGNAL             (1)
diff --git a/bricks/ev3rt/mpconfigport.h b/bricks/ev3rt/mpconfigport.h
@@ -43,6 +43,7 @@
 #define PYBRICKS_PY_ROBOTICS            (0)
 #define PYBRICKS_PY_ROBOTICS_DRIVEBASE_SPIKE (0)
 #define PYBRICKS_PY_TOOLS               (1)
+#define PYBRICKS_PY_TOOLS_HUB_MENU      (0)
 
 // Pybricks options
 #define PYBRICKS_OPT_COMPILER                   (1)
diff --git a/bricks/movehub/mpconfigport.h b/bricks/movehub/mpconfigport.h
@@ -37,6 +37,7 @@
 #define PYBRICKS_PY_ROBOTICS_DRIVEBASE_GYRO     (0)
 #define PYBRICKS_PY_ROBOTICS_DRIVEBASE_SPIKE    (0)
 #define PYBRICKS_PY_TOOLS                       (1)
+#define PYBRICKS_PY_TOOLS_HUB_MENU              (0)
 
 // Pybricks options
 #define PYBRICKS_OPT_COMPILER                   (0)
diff --git a/bricks/nxt/mpconfigport.h b/bricks/nxt/mpconfigport.h
@@ -42,6 +42,7 @@
 #define PYBRICKS_PY_ROBOTICS_DRIVEBASE_GYRO     (0)
 #define PYBRICKS_PY_ROBOTICS_DRIVEBASE_SPIKE    (0)
 #define PYBRICKS_PY_TOOLS                       (1)
+#define PYBRICKS_PY_TOOLS_HUB_MENU              (0)
 
 // Pybricks options
 #define PYBRICKS_OPT_COMPILER                   (1)
diff --git a/bricks/primehub/mpconfigport.h b/bricks/primehub/mpconfigport.h
@@ -43,6 +43,7 @@
 #define PYBRICKS_PY_ROBOTICS_DRIVEBASE_GYRO     (1)
 #define PYBRICKS_PY_ROBOTICS_DRIVEBASE_SPIKE    (1)
 #define PYBRICKS_PY_TOOLS                       (1)
+#define PYBRICKS_PY_TOOLS_HUB_MENU              (1)
 
 // Pybricks options
 #define PYBRICKS_OPT_COMPILER                   (1)
diff --git a/bricks/technichub/mpconfigport.h b/bricks/technichub/mpconfigport.h
@@ -41,6 +41,7 @@
 #define PYBRICKS_PY_ROBOTICS_DRIVEBASE_GYRO     (1)
 #define PYBRICKS_PY_ROBOTICS_DRIVEBASE_SPIKE    (0)
 #define PYBRICKS_PY_TOOLS                       (1)
+#define PYBRICKS_PY_TOOLS_HUB_MENU              (0)
 
 // Pybricks options
 #define PYBRICKS_OPT_COMPILER                   (1)
diff --git a/bricks/virtualhub/mpconfigvariant.h b/bricks/virtualhub/mpconfigvariant.h
@@ -36,6 +36,7 @@
 #define PYBRICKS_PY_ROBOTICS            (1)
 #define PYBRICKS_PY_ROBOTICS_DRIVEBASE_SPIKE (0)
 #define PYBRICKS_PY_TOOLS               (1)
+#define PYBRICKS_PY_TOOLS_HUB_MENU      (0)
 
 // Pybricks options
 #define PYBRICKS_OPT_COMPILER                   (1)
diff --git a/pybricks/tools/pb_module_tools.c b/pybricks/tools/pb_module_tools.c
@@ -13,6 +13,8 @@
 #include "py/stream.h"
 
 #include <pbio/int_math.h>
+#include <pbsys/light.h>
+#include <pbsys/program_stop.h>
 
 #include <pybricks/parameters.h>
 #include <pybricks/common.h>
@@ -143,10 +145,107 @@ STATIC mp_obj_t pb_module_tools_run_task(size_t n_args, const mp_obj_t *pos_args
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_KW(pb_module_tools_run_task_obj, 1, pb_module_tools_run_task);
 
+#if PYBRICKS_PY_TOOLS_HUB_MENU
+
+STATIC void pb_module_tools_hub_menu_display_symbol(mp_obj_t symbol) {
+    if (mp_obj_is_str(symbol)) {
+        pb_type_LightMatrix_display_char(pbsys_hub_light_matrix, symbol);
+    } else {
+        pb_type_LightMatrix_display_number(pbsys_hub_light_matrix, symbol);
+    }
+}
+
+/**
+ * Waits for a button press or release.
+ *
+ * @param [in]  press   Choose @c true to wait for press or @c false to wait for release.
+ * @returns             When waiting for pressed, it returns the button that was pressed, otherwise returns 0.
+ */
+STATIC pbio_button_flags_t pb_module_tools_hub_menu_wait_for_press(bool press) {
+
+    // This function should only be used in a blocking context.
+    pb_module_tools_assert_blocking();
+
+    pbio_error_t err;
+    pbio_button_flags_t btn;
+    while ((err = pbio_button_is_pressed(&btn)) == PBIO_SUCCESS && ((bool)btn) == !press) {
+        MICROPY_EVENT_POLL_HOOK;
+    }
+    pb_assert(err);
+    return btn;
+}
+
+/**
+ * Displays a menu on the hub display and allows the user to pick a symbol
+ * using the buttons.
+ *
+ * @param [in]  n_args  The number of args.
+ * @param [in]  args    The args passed in Python code (the menu entries).
+ */
+STATIC mp_obj_t pb_module_tools_hub_menu(size_t n_args, const mp_obj_t *args) {
+
+    // Validate arguments by displaying all of them, ending with the first.
+    // This ensures we fail right away instead of midway through the menu. It
+    // happens so fast that there isn't a time penalty for this.
+    for (int i = n_args - 1; i >= 0; i--) {
+        pb_module_tools_hub_menu_display_symbol(args[i]);
+    }
+
+    // Disable stop button and cache original setting to restore later.
+    pbio_button_flags_t stop_button = pbsys_program_stop_get_buttons();
+    pbsys_program_stop_set_buttons(0);
+
+    nlr_buf_t nlr;
+    if (nlr_push(&nlr) == 0) {
+
+        size_t selection = 0;
+
+        while (true) {
+            pb_module_tools_hub_menu_wait_for_press(false);
+            pbio_button_flags_t btn = pb_module_tools_hub_menu_wait_for_press(true);
+
+            // Selection made, exit.
+            if (btn & PBIO_BUTTON_CENTER) {
+                break;
+            }
+
+            // Increment/decrement selection for left/right buttons.
+            if (btn & PBIO_BUTTON_RIGHT) {
+                selection = (selection + 1) % n_args;
+            } else if (btn & PBIO_BUTTON_LEFT) {
+                selection = selection == 0 ? n_args - 1 : selection - 1;
+            }
+
+            // Display current selection.
+            pb_module_tools_hub_menu_display_symbol(args[selection]);
+        }
+
+        // Wait for release before returning, just like starting a normal program.
+        pb_module_tools_hub_menu_wait_for_press(false);
+
+        // Restore stop button setting prior to starting menu.
+        pbsys_program_stop_set_buttons(stop_button);
+
+        // Complete and return selected object.
+        nlr_pop();
+        return args[selection];
+    } else {
+        pbsys_program_stop_set_buttons(stop_button);
+        nlr_jump(nlr.ret_val);
+        return mp_const_none;
+    }
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(pb_module_tools_hub_menu_obj, 2, pb_module_tools_hub_menu);
+
+#endif // PYBRICKS_PY_TOOLS_HUB_MENU
+
 STATIC const mp_rom_map_elem_t tools_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR___name__),    MP_ROM_QSTR(MP_QSTR_tools)                    },
     { MP_ROM_QSTR(MP_QSTR_wait),        MP_ROM_PTR(&pb_module_tools_wait_obj)         },
     { MP_ROM_QSTR(MP_QSTR_read_input_byte), MP_ROM_PTR(&pb_module_tools_read_input_byte_obj) },
+    #if PYBRICKS_PY_TOOLS_HUB_MENU
+    { MP_ROM_QSTR(MP_QSTR_hub_menu),    MP_ROM_PTR(&pb_module_tools_hub_menu_obj)     },
+    #endif // PYBRICKS_PY_TOOLS_HUB_MENU
     { MP_ROM_QSTR(MP_QSTR_run_task),    MP_ROM_PTR(&pb_module_tools_run_task_obj)     },
     { MP_ROM_QSTR(MP_QSTR_StopWatch),   MP_ROM_PTR(&pb_type_StopWatch)                },
     { MP_ROM_QSTR(MP_QSTR_multitask),   MP_ROM_PTR(&pb_type_Task)                     },
