Skip to content

Commit

Permalink
style(python): improve type hints for input flows
Browse files Browse the repository at this point in the history
  • Loading branch information
matejcik committed Jan 13, 2025
1 parent de0aa25 commit e016e08
Show file tree
Hide file tree
Showing 3 changed files with 9 additions and 6 deletions.
9 changes: 6 additions & 3 deletions python/src/trezorlib/debuglink.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,8 @@ def __call__(
wait: bool | None = None,
) -> "LayoutContent": ...

InputFlowType = Generator[None, messages.ButtonRequest, None]


EXPECTED_RESPONSES_CONTEXT_LINES = 3

Expand Down Expand Up @@ -1108,7 +1110,7 @@ def _filter_message(self, msg: protobuf.MessageType) -> protobuf.MessageType:
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.
Expand Down Expand Up @@ -1142,7 +1144,7 @@ def set_input_flow(
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.
Expand Down Expand Up @@ -1190,7 +1192,8 @@ def __exit__(self, exc_type: Any, value: Any, traceback: Any) -> None:
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.
Expand Down
4 changes: 2 additions & 2 deletions tests/input_flows.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
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
Expand Down Expand Up @@ -2049,7 +2049,7 @@ def input_flow_common(self) -> BRGeneratorType:


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
Expand Down
2 changes: 1 addition & 1 deletion tests/input_flows_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -301,7 +301,7 @@ def input_mnemonic(self, mnemonic: list[str]) -> BRGeneratorType:

def input_all_slip39_shares(
self,
shares: list[str],
shares: t.Sequence[str],
has_groups: bool = False,
click_info: bool = False,
) -> BRGeneratorType:
Expand Down

0 comments on commit e016e08

Please sign in to comment.