[mypy] #163: Add type annotation on crypto_tools

silkaj/crypto_tools.py

@@ -15,6 +15,7 @@
 
 import hashlib
 import re
+from typing import Any, Optional, Union
 
 import base58
 
@@ -27,7 +28,7 @@ CHECKSUM_PATTERN = f"[1-9A-HJ-NP-Za-km-z]{{{CHECKSUM_SIZE}}}"
 PUBKEY_CHECKSUM_PATTERN = f"^{PUBKEY_PATTERN}:{CHECKSUM_PATTERN}$"
 
 
-def is_pubkey_and_check(pubkey):
+def is_pubkey_and_check(pubkey: str) -> Union[str, bool]:
     """
     Checks if the given argument contains a pubkey.
     If so, verifies the checksum if needed and returns the pubkey.
@@ -41,7 +42,7 @@ def is_pubkey_and_check(pubkey):
     return False
 
 
-def check_pubkey_format(pubkey, display_error=True):
+def check_pubkey_format(pubkey: str, display_error: bool = True) -> Optional[bool]:
     """
     Checks if a pubkey has a checksum.
     Exits if the pubkey is invalid.
@@ -52,10 +53,10 @@ def check_pubkey_format(pubkey, display_error=True):
         return True
     elif display_error:
         message_exit(f"Error: bad format for following public key: {pubkey}")
-    return
+    return None
 
 
-def validate_checksum(pubkey_checksum):
+def validate_checksum(pubkey_checksum: str) -> Any:
     """
     Check pubkey checksum after the pubkey, delimited by ":".
     If check pass: return pubkey
@@ -68,12 +69,13 @@ def validate_checksum(pubkey_checksum):
         f"Error: public key '{pubkey}' does not match checksum '{checksum}'.\n\
 Please verify the public key."
     )
+    return None
 
 
-def gen_checksum(pubkey):
+def gen_checksum(pubkey: str) -> str:
     """
     Returns the checksum of the input pubkey (encoded in b58)
     """
     pubkey_byte = base58.b58decode(pubkey)
     hash = hashlib.sha256(hashlib.sha256(pubkey_byte).digest()).digest()
-    return base58.b58encode(hash)[:3].decode("utf-8")
+    return str(base58.b58encode(hash)[:3].decode("utf-8"))