From 6506b02e2ec891a0a7fcca1d594d6adcc9618944 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Mon, 2 Sep 2024 12:59:56 +0200
Subject: [PATCH] feat(python/debuglink): introduce layout type awareness

[no changelog]
---
 python/src/trezorlib/debuglink.py | 30 +++++++++++++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

diff --git a/python/src/trezorlib/debuglink.py b/python/src/trezorlib/debuglink.py
index e5326a2df..843bfca6d 100644
--- a/python/src/trezorlib/debuglink.py
+++ b/python/src/trezorlib/debuglink.py
@@ -21,7 +21,7 @@ import textwrap
 import time
 from copy import deepcopy
 from datetime import datetime
-from enum import IntEnum
+from enum import Enum, IntEnum, auto
 from itertools import zip_longest
 from pathlib import Path
 from typing import (
@@ -65,6 +65,25 @@ EXPECTED_RESPONSES_CONTEXT_LINES = 3
 LOG = logging.getLogger(__name__)
 
 
+class LayoutType(Enum):
+    T1 = auto()
+    TT = auto()
+    TR = auto()
+    Mercury = auto()
+
+    @classmethod
+    def from_model(cls, model: models.TrezorModel) -> "LayoutType":
+        if model in (models.T2T1,):
+            return cls.TT
+        if model in (models.T2B1, models.T3B1):
+            return cls.TR
+        if model in (models.T3T1,):
+            return cls.Mercury
+        if model in (models.T1B1,):
+            return cls.T1
+        raise ValueError(f"Unknown model: {model}")
+
+
 class UnstructuredJSONReader:
     """Contains data-parsing helpers for JSON data that have unknown structure."""
 
@@ -404,6 +423,11 @@ class DebugLink:
         """Differences in handling debug events and LayoutContent."""
         return self.version < (2, 6, 1)
 
+    @property
+    def layout_type(self) -> LayoutType:
+        assert self.model is not None
+        return LayoutType.from_model(self.model)
+
     def set_screen_text_file(self, file_path: Optional[Path]) -> None:
         if file_path is not None:
             file_path.write_bytes(b"")
@@ -995,6 +1019,10 @@ class TrezorClientDebugLink(TrezorClient):
         self.debug.model = self.model
         self.debug.version = self.version
 
+    @property
+    def layout_type(self) -> LayoutType:
+        return self.debug.layout_type
+
     def reset_debug_features(self) -> None:
         """Prepare the debugging client for a new testcase.