From 0209768ff2d9ed448b786631f4af7424f4d0a19f Mon Sep 17 00:00:00 2001 From: matejcik Date: Thu, 30 Apr 2020 15:27:37 +0200 Subject: [PATCH] core/sign_tx: improve type hints --- core/src/apps/wallet/sign_tx/__init__.py | 2 +- core/src/apps/wallet/sign_tx/addresses.py | 6 ++++-- core/src/apps/wallet/sign_tx/helpers.py | 6 ++++-- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/core/src/apps/wallet/sign_tx/__init__.py b/core/src/apps/wallet/sign_tx/__init__.py index 407a2a2ede..6253bd0890 100644 --- a/core/src/apps/wallet/sign_tx/__init__.py +++ b/core/src/apps/wallet/sign_tx/__init__.py @@ -48,7 +48,7 @@ async def sign_tx(ctx: wire.Context, msg: SignTx, keychain: seed.Keychain) -> Tx except common.SigningError as e: raise wire.Error(*e.args) - res = None # type: Union[TxAck, bool] + res = None # type: Union[TxAck, bool, None] while True: try: req = signer.send(res) diff --git a/core/src/apps/wallet/sign_tx/addresses.py b/core/src/apps/wallet/sign_tx/addresses.py index de6bf89edc..522611ceb3 100644 --- a/core/src/apps/wallet/sign_tx/addresses.py +++ b/core/src/apps/wallet/sign_tx/addresses.py @@ -28,7 +28,7 @@ class AddressError(Exception): def get_address( - script_type: int, + script_type: EnumTypeInputScriptType, coin: CoinInfo, node: bip32.HDNode, multisig: MultisigRedeemScriptType = None, @@ -253,7 +253,9 @@ def validate_purpose(purpose: int, coin: CoinInfo) -> bool: return True -def validate_purpose_against_script_type(purpose: int, script_type: int) -> bool: +def validate_purpose_against_script_type( + purpose: int, script_type: EnumTypeInputScriptType +) -> bool: """ Validates purpose against provided input's script type: - 44 for spending address (script_type == SPENDADDRESS) diff --git a/core/src/apps/wallet/sign_tx/helpers.py b/core/src/apps/wallet/sign_tx/helpers.py index d619909ff8..520f09867a 100644 --- a/core/src/apps/wallet/sign_tx/helpers.py +++ b/core/src/apps/wallet/sign_tx/helpers.py @@ -22,7 +22,9 @@ from .writers import TX_HASH_SIZE from apps.common.coininfo import CoinInfo if False: - from typing import Any, Awaitable, Union + from typing import Any, Awaitable, Dict, Union + from trezor.messages.TxInputType import EnumTypeInputScriptType + from trezor.messages.TxOutputType import EnumTypeOutputScriptType MULTISIG_INPUT_SCRIPT_TYPES = ( InputScriptType.SPENDMULTISIG, @@ -40,7 +42,7 @@ CHANGE_OUTPUT_TO_INPUT_SCRIPT_TYPES = { OutputScriptType.PAYTOMULTISIG: InputScriptType.SPENDMULTISIG, OutputScriptType.PAYTOP2SHWITNESS: InputScriptType.SPENDP2SHWITNESS, OutputScriptType.PAYTOWITNESS: InputScriptType.SPENDWITNESS, -} +} # type: Dict[EnumTypeOutputScriptType, EnumTypeInputScriptType] INTERNAL_INPUT_SCRIPT_TYPES = tuple(CHANGE_OUTPUT_TO_INPUT_SCRIPT_TYPES.values()) CHANGE_OUTPUT_SCRIPT_TYPES = tuple(CHANGE_OUTPUT_TO_INPUT_SCRIPT_TYPES.keys())