From 6feada2eed189f8fd7aa7e3a685871f10aaaaad2 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Wed, 13 Apr 2022 15:42:28 +0200
Subject: [PATCH] style(python/monero): improve type signature on Monero
 functions

---
 python/.changelog.d/2219.changed   |  1 +
 python/src/trezorlib/cli/monero.py | 25 +++++++++++++++++--------
 python/src/trezorlib/monero.py     |  6 ++++--
 3 files changed, 22 insertions(+), 10 deletions(-)
 create mode 100644 python/.changelog.d/2219.changed

diff --git a/python/.changelog.d/2219.changed b/python/.changelog.d/2219.changed
new file mode 100644
index 000000000..6aa229ee3
--- /dev/null
+++ b/python/.changelog.d/2219.changed
@@ -0,0 +1 @@
+`trezorctl monero` network type arguments now accept symbolic names instead of numbers.
diff --git a/python/src/trezorlib/cli/monero.py b/python/src/trezorlib/cli/monero.py
index 451c2b30d..d06fbc165 100644
--- a/python/src/trezorlib/cli/monero.py
+++ b/python/src/trezorlib/cli/monero.py
@@ -18,8 +18,8 @@ from typing import TYPE_CHECKING, Dict
 
 import click
 
-from .. import monero, tools
-from . import with_client
+from .. import messages, monero, tools
+from . import ChoiceType, with_client
 
 if TYPE_CHECKING:
     from ..client import TrezorClient
@@ -36,29 +36,38 @@ def cli() -> None:
 @click.option("-n", "--address", required=True, help=PATH_HELP)
 @click.option("-d", "--show-display", is_flag=True)
 @click.option(
-    "-t", "--network-type", type=click.Choice(["0", "1", "2", "3"]), default="0"
+    "-t",
+    "--network-type",
+    type=ChoiceType({m.name: m for m in messages.MoneroNetworkType}),
+    default=messages.MoneroNetworkType.MAINNET,
 )
 @with_client
 def get_address(
-    client: "TrezorClient", address: str, show_display: bool, network_type: str
+    client: "TrezorClient",
+    address: str,
+    show_display: bool,
+    network_type: messages.MoneroNetworkType,
 ) -> bytes:
     """Get Monero address for specified path."""
     address_n = tools.parse_path(address)
-    return monero.get_address(client, address_n, show_display, int(network_type))
+    return monero.get_address(client, address_n, show_display, network_type)
 
 
 @cli.command()
 @click.option("-n", "--address", required=True, help=PATH_HELP)
 @click.option(
-    "-t", "--network-type", type=click.Choice(["0", "1", "2", "3"]), default="0"
+    "-t",
+    "--network-type",
+    type=ChoiceType({m.name: m for m in messages.MoneroNetworkType}),
+    default=messages.MoneroNetworkType.MAINNET,
 )
 @with_client
 def get_watch_key(
-    client: "TrezorClient", address: str, network_type: str
+    client: "TrezorClient", address: str, network_type: messages.MoneroNetworkType
 ) -> Dict[str, str]:
     """Get Monero watch key for specified path."""
     address_n = tools.parse_path(address)
-    res = monero.get_watch_key(client, address_n, int(network_type))
+    res = monero.get_watch_key(client, address_n, network_type)
     # TODO: could be made required in MoneroWatchKey
     assert res.address is not None
     assert res.watch_key is not None
diff --git a/python/src/trezorlib/monero.py b/python/src/trezorlib/monero.py
index 4414b7ab3..4b9187f16 100644
--- a/python/src/trezorlib/monero.py
+++ b/python/src/trezorlib/monero.py
@@ -36,7 +36,7 @@ def get_address(
     client: "TrezorClient",
     n: "Address",
     show_display: bool = False,
-    network_type: int = 0,
+    network_type: messages.MoneroNetworkType = messages.MoneroNetworkType.MAINNET,
 ) -> "MessageType":
     return client.call(
         messages.MoneroGetAddress(
@@ -47,7 +47,9 @@ def get_address(
 
 @expect(messages.MoneroWatchKey)
 def get_watch_key(
-    client: "TrezorClient", n: "Address", network_type: int = 0
+    client: "TrezorClient",
+    n: "Address",
+    network_type: messages.MoneroNetworkType = messages.MoneroNetworkType.MAINNET,
 ) -> "MessageType":
     return client.call(
         messages.MoneroGetWatchKey(address_n=n, network_type=network_type)