diff --git a/python/src/trezorlib/_internal/emulator.py b/python/src/trezorlib/_internal/emulator.py index 39704803c..f1c3ed9fd 100644 --- a/python/src/trezorlib/_internal/emulator.py +++ b/python/src/trezorlib/_internal/emulator.py @@ -72,7 +72,8 @@ class Emulator: else: self.logfile = self.profile_dir / "trezor.log" - self.client: Optional[TrezorClientDebugLink] = None + # Using `client` property instead to assert `not None` + self._client: Optional[TrezorClientDebugLink] = None self.process: Optional[subprocess.Popen] = None self.port = 21324 @@ -81,6 +82,16 @@ class Emulator: self.auto_interact = auto_interact self.extra_args = list(extra_args) + @property + def client(self) -> TrezorClientDebugLink: + """So that type-checkers do not see `client` as `Optional`. + + (it is not None between `start()` and `stop()` calls) + """ + if self._client is None: + raise RuntimeError + return self._client + def make_args(self) -> List[str]: return [] @@ -162,14 +173,15 @@ class Emulator: (self.profile_dir / "trezor.port").write_text(str(self.port) + "\n") transport = self._get_transport() - self.client = TrezorClientDebugLink(transport, auto_interact=self.auto_interact) - - self.client.open() + self._client = TrezorClientDebugLink( + transport, auto_interact=self.auto_interact + ) + self._client.open() def stop(self) -> None: - if self.client: - self.client.close() - self.client = None + if self._client: + self._client.close() + self._client = None if self.process: LOG.info("Terminating emulator...") diff --git a/python/src/trezorlib/debuglink.py b/python/src/trezorlib/debuglink.py index 57a9bfbe0..a05c38809 100644 --- a/python/src/trezorlib/debuglink.py +++ b/python/src/trezorlib/debuglink.py @@ -36,9 +36,11 @@ from typing import ( Tuple, Type, Union, + overload, ) from mnemonic import Mnemonic +from typing_extensions import Literal from . import mapping, messages, protobuf from .client import TrezorClient @@ -296,6 +298,17 @@ class DebugLink: return None + # Type overloads make sure that when we supply `wait=True` into `click()`, + # it will always return `LayoutContent` and we do not need to assert `is not None`. + + @overload + def click(self, click: Tuple[int, int]) -> None: + ... + + @overload + def click(self, click: Tuple[int, int], wait: Literal[True]) -> LayoutContent: + ... + def click( self, click: Tuple[int, int], wait: bool = False ) -> Optional[LayoutContent]: