From c4bf4fa884f3ba61b543654aadae345c0863cca3 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Fri, 17 Feb 2023 11:44:31 +0100
Subject: [PATCH] fix(python): improve type checking

---
 python/.changelog.d/2832.fixed           | 1 +
 python/src/trezorlib/client.py           | 8 +++++---
 python/src/trezorlib/transport/bridge.py | 2 +-
 python/src/trezorlib/transport/udp.py    | 2 +-
 4 files changed, 8 insertions(+), 5 deletions(-)
 create mode 100644 python/.changelog.d/2832.fixed

diff --git a/python/.changelog.d/2832.fixed b/python/.changelog.d/2832.fixed
new file mode 100644
index 000000000..ae0edb417
--- /dev/null
+++ b/python/.changelog.d/2832.fixed
@@ -0,0 +1 @@
+Improve typing information when `TrezorClient` has a more intelligent UI object.
diff --git a/python/src/trezorlib/client.py b/python/src/trezorlib/client.py
index 684093b03..85bcaa340 100644
--- a/python/src/trezorlib/client.py
+++ b/python/src/trezorlib/client.py
@@ -17,7 +17,7 @@
 import logging
 import os
 import warnings
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar
 
 from mnemonic import Mnemonic
 
@@ -31,6 +31,8 @@ if TYPE_CHECKING:
     from .ui import TrezorClientUI
     from .transport import Transport
 
+UI = TypeVar("UI", bound="TrezorClientUI")
+
 LOG = logging.getLogger(__name__)
 
 MAX_PASSPHRASE_LENGTH = 50
@@ -70,7 +72,7 @@ def get_default_client(
     return TrezorClient(transport, ui, **kwargs)
 
 
-class TrezorClient:
+class TrezorClient(Generic[UI]):
     """Trezor client, a connection to a Trezor device.
 
     This class allows you to manage connection state, send and receive protobuf
@@ -81,7 +83,7 @@ class TrezorClient:
     def __init__(
         self,
         transport: "Transport",
-        ui: "TrezorClientUI",
+        ui: UI,
         session_id: Optional[bytes] = None,
         derive_cardano: Optional[bool] = None,
         model: Optional[models.TrezorModel] = None,
diff --git a/python/src/trezorlib/transport/bridge.py b/python/src/trezorlib/transport/bridge.py
index 0066e12d0..e0c34a8f7 100644
--- a/python/src/trezorlib/transport/bridge.py
+++ b/python/src/trezorlib/transport/bridge.py
@@ -109,7 +109,7 @@ class BridgeTransport(Transport):
     """
 
     PATH_PREFIX = "bridge"
-    ENABLED = True
+    ENABLED: bool = True
 
     def __init__(
         self, device: Dict[str, Any], legacy: bool, debug: bool = False
diff --git a/python/src/trezorlib/transport/udp.py b/python/src/trezorlib/transport/udp.py
index c79df9350..f305a77c2 100644
--- a/python/src/trezorlib/transport/udp.py
+++ b/python/src/trezorlib/transport/udp.py
@@ -36,7 +36,7 @@ class UdpTransport(ProtocolBasedTransport):
     DEFAULT_HOST = "127.0.0.1"
     DEFAULT_PORT = 21324
     PATH_PREFIX = "udp"
-    ENABLED = True
+    ENABLED: bool = True
 
     def __init__(self, device: Optional[str] = None) -> None:
         if not device: