refactor(core): return uPy Obj from DeviceMenu

[no changelog]

core/embed/rust/librust_qstr.h

@@ -18,12 +18,18 @@ static void _librust_qstrs(void) {
   MP_QSTR_;
   MP_QSTR_ATTACHED;
   MP_QSTR_AttachType;
+  MP_QSTR_AutoLockDelay;
   MP_QSTR_BacklightLevels;
+  MP_QSTR_BackupFailed;
   MP_QSTR_BleInterface;
   MP_QSTR_CANCELLED;
   MP_QSTR_CONFIRMED;
+  MP_QSTR_CheckBackup;
   MP_QSTR_DIM;
   MP_QSTR_DONE;
+  MP_QSTR_DeviceDisconnect;
+  MP_QSTR_DeviceMenuResult;
+  MP_QSTR_DevicePair;
   MP_QSTR_INFO;
   MP_QSTR_INITIAL;
   MP_QSTR_LOW;
@@ -42,10 +48,12 @@ static void _librust_qstrs(void) {
   MP_QSTR_SWIPE_LEFT;
   MP_QSTR_SWIPE_RIGHT;
   MP_QSTR_SWIPE_UP;
+  MP_QSTR_ScreenBrightness;
   MP_QSTR_TR;
   MP_QSTR_TRANSITIONING;
   MP_QSTR_TX_PACKET_LEN;
   MP_QSTR_TranslationsHeader;
+  MP_QSTR_WipeDevice;
   MP_QSTR___del__;
   MP_QSTR___dict__;
   MP_QSTR___name__;

core/embed/rust/src/ui/api/firmware_micropython.rs

@@ -17,6 +17,7 @@ use crate::{
         component::Empty,
         layout::{
             base::LAYOUT_STATE,
+            device_menu_result::DEVICE_MENU_RESULT,
             obj::{ComponentMsgObj, LayoutObj, ATTACH_TYPE_OBJ},
             result::{CANCELLED, CONFIRMED, INFO},
             util::{upy_disable_animation, RecoveryType},
@@ -1665,7 +1666,7 @@ pub static mp_module_trezorui_api: Module = obj_module! {
     ///     device_name: str,
     ///     paired_devices: Iterable[str],
     ///     auto_lock_delay: str,
-    /// ) -> LayoutObj[UiResult]:
+    /// ) -> LayoutObj[UiResult | DeviceMenuResult | tuple[DeviceMenuResult, int]]:
     ///     """Show the device menu."""
     Qstr::MP_QSTR_show_device_menu => obj_fn_kw!(0, new_show_device_menu).as_obj(),
 
@@ -1834,4 +1835,14 @@ pub static mp_module_trezorui_api: Module = obj_module! {
     ///     DONE: "ClassVar[LayoutState]"
     Qstr::MP_QSTR_LayoutState => LAYOUT_STATE.as_obj(),
 
+    /// class DeviceMenuResult:
+    ///     """Result of a device menu operation."""
+    ///     BackupFailed: ClassVar[DeviceMenuResult]
+    ///     DevicePair: ClassVar[DeviceMenuResult]
+    ///     DeviceDisconnect: ClassVar[DeviceMenuResult]
+    ///     CheckBackup: ClassVar[DeviceMenuResult]
+    ///     WipeDevice: ClassVar[DeviceMenuResult]
+    ///     ScreenBrightness: ClassVar[DeviceMenuResult]
+    ///     AutoLockDelay: ClassVar[DeviceMenuResult]
+    Qstr::MP_QSTR_DeviceMenuResult => DEVICE_MENU_RESULT.as_obj(),
 };

core/embed/rust/src/ui/layout/device_menu_result.rs

@@ -0,0 +1,32 @@
+use crate::micropython::{
+    macros::{obj_dict, obj_map, obj_type},
+    qstr::Qstr,
+    simple_type::SimpleTypeObj,
+    typ::Type,
+};
+
+static DEVICE_MENU_RESULT_BASE_TYPE: Type = obj_type! { name: Qstr::MP_QSTR_DeviceMenuResult, };
+
+pub static BACKUP_FAILED: SimpleTypeObj = SimpleTypeObj::new(&DEVICE_MENU_RESULT_BASE_TYPE);
+pub static DEVICE_PAIR: SimpleTypeObj = SimpleTypeObj::new(&DEVICE_MENU_RESULT_BASE_TYPE);
+pub static DEVICE_DISCONNECT: SimpleTypeObj = SimpleTypeObj::new(&DEVICE_MENU_RESULT_BASE_TYPE);
+pub static CHECK_BACKUP: SimpleTypeObj = SimpleTypeObj::new(&DEVICE_MENU_RESULT_BASE_TYPE);
+pub static WIPE_DEVICE: SimpleTypeObj = SimpleTypeObj::new(&DEVICE_MENU_RESULT_BASE_TYPE);
+pub static SCREEN_BRIGHTNESS: SimpleTypeObj = SimpleTypeObj::new(&DEVICE_MENU_RESULT_BASE_TYPE);
+pub static AUTO_LOCK_DELAY: SimpleTypeObj = SimpleTypeObj::new(&DEVICE_MENU_RESULT_BASE_TYPE);
+
+// Create a DeviceMenuResult class that contains all result types
+static DEVICE_MENU_RESULT_TYPE: Type = obj_type! {
+    name: Qstr::MP_QSTR_DeviceMenuResult,
+    locals: &obj_dict! { obj_map! {
+        Qstr::MP_QSTR_BackupFailed => BACKUP_FAILED.as_obj(),
+        Qstr::MP_QSTR_DevicePair => DEVICE_PAIR.as_obj(),
+        Qstr::MP_QSTR_DeviceDisconnect => DEVICE_DISCONNECT.as_obj(),
+        Qstr::MP_QSTR_CheckBackup => CHECK_BACKUP.as_obj(),
+        Qstr::MP_QSTR_WipeDevice => WIPE_DEVICE.as_obj(),
+        Qstr::MP_QSTR_ScreenBrightness => SCREEN_BRIGHTNESS.as_obj(),
+        Qstr::MP_QSTR_AutoLockDelay => AUTO_LOCK_DELAY.as_obj(),
+    } },
+};
+
+pub static DEVICE_MENU_RESULT: SimpleTypeObj = SimpleTypeObj::new(&DEVICE_MENU_RESULT_TYPE);

core/embed/rust/src/ui/layout/mod.rs

@@ -3,6 +3,8 @@ pub mod base;
 #[cfg(feature = "micropython")]
 pub mod obj;
 
+#[cfg(feature = "micropython")]
+pub mod device_menu_result;
 #[cfg(feature = "micropython")]
 pub mod result;
 

core/embed/rust/src/ui/layout_eckhart/component_msg_obj.rs

@@ -1,12 +1,13 @@
 use crate::{
     error::Error,
-    micropython::obj::Obj,
+    micropython::{obj::Obj, util::new_tuple},
     ui::{
         component::{
             text::paragraphs::{ParagraphSource, Paragraphs},
             Component, Timeout,
         },
         layout::{
+            device_menu_result::*,
             obj::ComponentMsgObj,
             result::{CANCELLED, CONFIRMED, INFO},
         },
@@ -148,13 +149,15 @@ impl ComponentMsgObj for SetBrightnessScreen {
 impl<'a> ComponentMsgObj for DeviceMenuScreen<'a> {
     fn msg_try_into_obj(&self, msg: Self::Msg) -> Result<Obj, Error> {
         match msg {
-            DeviceMenuMsg::BackupFailed => "BackupFailed".try_into(),
-            DeviceMenuMsg::DevicePair => "DevicePair".try_into(),
-            DeviceMenuMsg::DeviceDisconnect(_) => "DeviceDisconnect".try_into(),
-            DeviceMenuMsg::CheckBackup => "CheckBackup".try_into(),
-            DeviceMenuMsg::WipeDevice => "WipeDevice".try_into(),
-            DeviceMenuMsg::ScreenBrightness => "ScreenBrightness".try_into(),
-            DeviceMenuMsg::AutoLockDelay => "AutoLockDelay".try_into(),
+            DeviceMenuMsg::BackupFailed => Ok(BACKUP_FAILED.as_obj()),
+            DeviceMenuMsg::DevicePair => Ok(DEVICE_PAIR.as_obj()),
+            DeviceMenuMsg::DeviceDisconnect(index) => {
+                Ok(new_tuple(&[DEVICE_DISCONNECT.as_obj(), index.try_into()?])?)
+            }
+            DeviceMenuMsg::CheckBackup => Ok(CHECK_BACKUP.as_obj()),
+            DeviceMenuMsg::WipeDevice => Ok(WIPE_DEVICE.as_obj()),
+            DeviceMenuMsg::ScreenBrightness => Ok(SCREEN_BRIGHTNESS.as_obj()),
+            DeviceMenuMsg::AutoLockDelay => Ok(AUTO_LOCK_DELAY.as_obj()),
             DeviceMenuMsg::Close => Ok(CANCELLED.as_obj()),
         }
     }

core/mocks/generated/trezorui_api.pyi

@@ -555,7 +555,7 @@ def show_device_menu(
     device_name: str,
     paired_devices: Iterable[str],
     auto_lock_delay: str,
-) -> LayoutObj[UiResult]:
+) -> LayoutObj[UiResult | DeviceMenuResult | tuple[DeviceMenuResult, int]]:
     """Show the device menu."""
 
 
@@ -739,3 +739,15 @@ class LayoutState:
     ATTACHED: "ClassVar[LayoutState]"
     TRANSITIONING: "ClassVar[LayoutState]"
     DONE: "ClassVar[LayoutState]"
+
+
+# rust/src/ui/api/firmware_micropython.rs
+class DeviceMenuResult:
+    """Result of a device menu operation."""
+    BackupFailed: ClassVar[DeviceMenuResult]
+    DevicePair: ClassVar[DeviceMenuResult]
+    DeviceDisconnect: ClassVar[DeviceMenuResult]
+    CheckBackup: ClassVar[DeviceMenuResult]
+    WipeDevice: ClassVar[DeviceMenuResult]
+    ScreenBrightness: ClassVar[DeviceMenuResult]
+    AutoLockDelay: ClassVar[DeviceMenuResult]

core/src/apps/homescreen/device_menu.py

@@ -1,9 +1,9 @@
 import storage.device
 import trezorui_api
 from trezor import TR, config, utils
-from trezor.ui.layouts import raise_if_not_confirmed
-from trezor.ui.layouts.common import interact
+from trezor.ui.layouts import interact
 from trezor.wire import ActionCancelled
+from trezorui_api import DeviceMenuResult
 
 
 async def _prompt_auto_lock_delay() -> int:
@@ -61,7 +61,7 @@ async def handle_device_menu() -> None:
         "{count} {plural}", auto_lock_num, auto_lock_label
     )
 
-    menu_result = await raise_if_not_confirmed(
+    menu_result = await interact(
         trezorui_api.show_device_menu(
             failed_backup=failed_backup,
             battery_percentage=battery_percentage,
@@ -70,27 +70,37 @@ async def handle_device_menu() -> None:
             device_name=device_name,
             auto_lock_delay=auto_lock_str,
         ),
-        None,
+        "device_menu",
     )
-    if menu_result == "DevicePair":
+
+    if menu_result is DeviceMenuResult.DevicePair:
         from apps.management.ble.pair_new_device import pair_new_device
 
         await pair_new_device()
-    elif menu_result == "ScreenBrightness":
+    elif menu_result is DeviceMenuResult.ScreenBrightness:
         from trezor.ui.layouts import set_brightness
 
         await set_brightness()
-    elif menu_result == "WipeDevice":
+    elif menu_result is DeviceMenuResult.WipeDevice:
         from trezor.messages import WipeDevice
 
         from apps.management.wipe_device import wipe_device
 
         await wipe_device(WipeDevice())
-    elif menu_result == "AutoLockDelay":
+    elif menu_result is DeviceMenuResult.AutoLockDelay:
 
         if config.has_pin():
 
             auto_lock_delay_ms = await _prompt_auto_lock_delay()
             storage.device.set_autolock_delay_ms(auto_lock_delay_ms)
+    elif isinstance(menu_result, tuple):
+        # It's a tuple with (result_type, index)
+        result_type, index = menu_result
+        if result_type is DeviceMenuResult.DeviceDisconnect:
+            raise RuntimeError(
+                f"Device disconnect not implemented, device index: {index}"
+            )
+        else:
+            raise RuntimeError(f"Unknown menu {result_type}, {index}")
     else:
         raise RuntimeError(f"Unknown menu {menu_result}")