Add Type Hints to Bech32 Functions

bit/base32.py

@@ -1,30 +1,9 @@
-# Copyright (c) 2017 Pieter Wuille
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in
-# all copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-# THE SOFTWARE.
-
-from bit.constants import BECH32_VERSION_SET
+from typing import List, Optional, Tuple
 
 CHARSET = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"
 
 
-def bech32_polymod(values):
-    """Internal function that computes the Bech32 checksum."""
+def bech32_polymod(values: List[int]) -> int:
     generator = [0x3B6A57B2, 0x26508E6D, 0x1EA119FA, 0x3D4233DD, 0x2A1462B3]
     chk = 1
     for value in values:
@@ -35,48 +14,42 @@ def bech32_polymod(values):
     return chk
 
 
-def bech32_hrp_expand(hrp):
-    """Expand the HRP into values for checksum computation."""
+def bech32_hrp_expand(hrp: str) -> List[int]:
     return [ord(x) >> 5 for x in hrp] + [0] + [ord(x) & 31 for x in hrp]
 
 
-def bech32_verify_checksum(hrp, data):
-    """Verify a checksum given HRP and converted data characters."""
+def bech32_verify_checksum(hrp: str, data: List[int]) -> bool:
     return bech32_polymod(bech32_hrp_expand(hrp) + data) == 1
 
 
-def bech32_create_checksum(hrp, data):
-    """Compute the checksum values given HRP and data."""
+def bech32_create_checksum(hrp: str, data: List[int]) -> List[int]:
     values = bech32_hrp_expand(hrp) + data
     polymod = bech32_polymod(values + [0, 0, 0, 0, 0, 0]) ^ 1
     return [(polymod >> 5 * (5 - i)) & 31 for i in range(6)]
 
 
-def bech32_encode(hrp, data):
-    """Compute a Bech32 string given HRP and data values."""
+def bech32_encode(hrp: str, data: List[int]) -> str:
     combined = data + bech32_create_checksum(hrp, data)
     return hrp + '1' + ''.join([CHARSET[d] for d in combined])
 
 
-def bech32_decode(bech):
-    """Validate a Bech32 string, and determine HRP and data."""
+def bech32_decode(bech: str) -> Tuple[Optional[str], Optional[List[int]]]:
     if (any(ord(x) < 33 or ord(x) > 126 for x in bech)) or (bech.lower() != bech and bech.upper() != bech):
         return (None, None)
     bech = bech.lower()
     pos = bech.rfind('1')
     if pos < 1 or pos + 7 > len(bech) or len(bech) > 90:
         return (None, None)
-    if not all(x in CHARSET for x in bech[pos + 1 :]):
+    if not all(x in CHARSET for x in bech[pos + 1:]):
         return (None, None)
     hrp = bech[:pos]
-    data = [CHARSET.find(x) for x in bech[pos + 1 :]]
+    data = [CHARSET.find(x) for x in bech[pos + 1:]]
     if not bech32_verify_checksum(hrp, data):
         return (None, None)
     return (hrp, data[:-6])
 
 
-def convertbits(data, frombits, tobits, pad=True):
-    """General power-of-2 base conversion."""
+def convertbits(data: List[int], frombits: int, tobits: int, pad: bool = True) -> Optional[List[int]]:
     acc = 0
     bits = 0
     ret = []
@@ -98,11 +71,8 @@ def convertbits(data, frombits, tobits, pad=True):
     return ret
 
 
-# def decode(hrp, addr):
-def decode(addr):
-    """Decode a segwit address."""
+def decode(addr: str) -> Tuple[Optional[int], Optional[List[int]]]:
     hrpgot, data = bech32_decode(addr)
-    # if hrpgot != hrp:
     if hrpgot not in BECH32_VERSION_SET:
         return (None, None)
     decoded = convertbits(data[1:], 5, 8, False)
@@ -115,8 +85,7 @@ def decode(addr):
     return (data[0], decoded)
 
 
-def encode(hrp, witver, witprog):
-    """Encode a segwit address."""
+def encode(hrp: str, witver: int, witprog: List[int]) -> Optional[str]:
     ret = bech32_encode(hrp, [witver] + convertbits(witprog, 8, 5))
     if decode(ret) == (None, None):
         return None