From 79cc7a96938e3bf16c9a0b72b53792f2dec4185c Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Thu, 9 Jan 2025 15:11:29 +0100
Subject: [PATCH] style(python): improve type hints for input flows

---
 python/src/trezorlib/debuglink.py | 9 ++++++---
 tests/input_flows.py              | 4 ++--
 tests/input_flows_helpers.py      | 2 +-
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/python/src/trezorlib/debuglink.py b/python/src/trezorlib/debuglink.py
index 56c1dfa344..3af5711a7c 100644
--- a/python/src/trezorlib/debuglink.py
+++ b/python/src/trezorlib/debuglink.py
@@ -67,6 +67,8 @@ if TYPE_CHECKING:
             wait: bool | None = None,
         ) -> "LayoutContent": ...
 
+    InputFlowType = Generator[None, messages.ButtonRequest, None]
+
 
 EXPECTED_RESPONSES_CONTEXT_LINES = 3
 
@@ -1108,7 +1110,7 @@ class TrezorClientDebugLink(TrezorClient):
             return msg
 
     def set_input_flow(
-        self, input_flow: Generator[None, messages.ButtonRequest | None, None]
+        self, input_flow: InputFlowType | Callable[[], InputFlowType]
     ) -> None:
         """Configure a sequence of input events for the current with-block.
 
@@ -1142,7 +1144,7 @@ class TrezorClientDebugLink(TrezorClient):
         if not hasattr(input_flow, "send"):
             raise RuntimeError("input_flow should be a generator function")
         self.ui.input_flow = input_flow
-        input_flow.send(None)  # start the generator
+        next(input_flow)  # start the generator
 
     def watch_layout(self, watch: bool = True) -> None:
         """Enable or disable watching layout changes.
@@ -1190,7 +1192,8 @@ class TrezorClientDebugLink(TrezorClient):
             input_flow.throw(exc_type, value, traceback)
 
     def set_expected_responses(
-        self, expected: list[Union["ExpectedMessage", Tuple[bool, "ExpectedMessage"]]]
+        self,
+        expected: Sequence[Union["ExpectedMessage", Tuple[bool, "ExpectedMessage"]]],
     ) -> None:
         """Set a sequence of expected responses to client calls.
 
diff --git a/tests/input_flows.py b/tests/input_flows.py
index 2609cada30..00bd63ea08 100644
--- a/tests/input_flows.py
+++ b/tests/input_flows.py
@@ -12,7 +12,7 @@ input flow details.
 from __future__ import annotations
 
 import time
-from typing import Callable, Generator
+from typing import Callable, Generator, Sequence
 
 from trezorlib import messages
 from trezorlib.debuglink import DebugLink, LayoutContent, LayoutType
@@ -2049,7 +2049,7 @@ class InputFlowSlip39BasicRecoveryDryRun(InputFlowBase):
 
 
 class InputFlowSlip39BasicRecovery(InputFlowBase):
-    def __init__(self, client: Client, shares: list[str], pin: str | None = None):
+    def __init__(self, client: Client, shares: Sequence[str], pin: str | None = None):
         super().__init__(client)
         self.shares = shares
         self.pin = pin
diff --git a/tests/input_flows_helpers.py b/tests/input_flows_helpers.py
index 3ad5f580ed..d51d3a8315 100644
--- a/tests/input_flows_helpers.py
+++ b/tests/input_flows_helpers.py
@@ -301,7 +301,7 @@ class RecoveryFlow:
 
     def input_all_slip39_shares(
         self,
-        shares: list[str],
+        shares: t.Sequence[str],
         has_groups: bool = False,
         click_info: bool = False,
     ) -> BRGeneratorType: