pybricks.tools: Add hub_menu function.

This implements the first approach discussed in pybricks/support#1064.
pybricks · Jun 28, 2023 · 0fdcc2d · 0fdcc2d
1 parent b5f40d8
commit 0fdcc2d
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 0 deletions.
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/pybricks/tools/pb_module_tools.c b/pybricks/tools/pb_module_tools.c
@@ -13,6 +13,7 @@
 #include "py/stream.h"
 
 #include <pbio/int_math.h>
+#include <pbsys/program_stop.h>
 
 #include <pybricks/parameters.h>
 #include <pybricks/common.h>
@@ -143,10 +144,88 @@ 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_COMMON && PYBRICKS_PY_COMMON_LIGHT_MATRIX
+
+#include <pbsys/light.h>
+
+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);
+    }
+}
+
+STATIC pbio_button_flags_t pb_module_tools_hub_menu_wait_for_press(bool press) {
+    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);
+
+    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);
+    return args[selection];
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(pb_module_tools_hub_menu_obj, 2, pb_module_tools_hub_menu);
+
+#endif // PYBRICKS_PY_COMMON && PYBRICKS_PY_COMMON_LIGHT_MATRIX
+
 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_COMMON && PYBRICKS_PY_COMMON_LIGHT_MATRIX
+    { MP_ROM_QSTR(MP_QSTR_hub_menu),    MP_ROM_PTR(&pb_module_tools_hub_menu_obj)     },
+    #endif // PYBRICKS_PY_COMMON && PYBRICKS_PY_COMMON_LIGHT_MATRIX
     { 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)                     },