Skip to content

Commit cafaad7

Browse files
cdeckerrustyrussell
authored andcommitted
pyln: Add type-annotations to pyln-proto
1 parent 2a176d2 commit cafaad7

File tree

3 files changed

+23
-18
lines changed

3 files changed

+23
-18
lines changed

contrib/pyln-proto/pyln/proto/bech32.py

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,13 @@
1919
# THE SOFTWARE.
2020

2121
"""Reference implementation for Bech32 and segwit addresses."""
22+
from typing import Tuple
2223

2324

2425
CHARSET = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"
2526

2627

27-
def bech32_polymod(values):
28+
def bech32_polymod(values: bytes) -> int:
2829
"""Internal function that computes the Bech32 checksum."""
2930
generator = [0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3]
3031
chk = 1
@@ -36,43 +37,47 @@ def bech32_polymod(values):
3637
return chk
3738

3839

39-
def bech32_hrp_expand(hrp):
40+
def bech32_hrp_expand(hrp: str) -> bytes:
4041
"""Expand the HRP into values for checksum computation."""
41-
return [ord(x) >> 5 for x in hrp] + [0] + [ord(x) & 31 for x in hrp]
42+
return bytes([ord(x) >> 5 for x in hrp] + [0] + [ord(x) & 31 for x in hrp])
4243

4344

44-
def bech32_verify_checksum(hrp, data):
45+
def bech32_verify_checksum(hrp: str, data: bytes) -> bool:
4546
"""Verify a checksum given HRP and converted data characters."""
4647
return bech32_polymod(bech32_hrp_expand(hrp) + data) == 1
4748

4849

49-
def bech32_create_checksum(hrp, data):
50+
def bech32_create_checksum(hrp: str, data: bytes) -> bytes:
5051
"""Compute the checksum values given HRP and data."""
5152
values = bech32_hrp_expand(hrp) + data
52-
polymod = bech32_polymod(values + [0, 0, 0, 0, 0, 0]) ^ 1
53-
return [(polymod >> 5 * (5 - i)) & 31 for i in range(6)]
53+
polymod = bech32_polymod(values + bytes([0, 0, 0, 0, 0, 0])) ^ 1
54+
return bytes([(polymod >> 5 * (5 - i)) & 31 for i in range(6)])
5455

5556

56-
def bech32_encode(hrp, data):
57+
def bech32_encode(hrp: str, data: bytes) -> str:
5758
"""Compute a Bech32 string given HRP and data values."""
5859
combined = data + bech32_create_checksum(hrp, data)
5960
return hrp + '1' + ''.join([CHARSET[d] for d in combined])
6061

6162

62-
def bech32_decode(bech):
63+
def bech32_decode(bech: str) -> Tuple[str, bytes]:
6364
"""Validate a Bech32 string, and determine HRP and data."""
6465
if ((any(ord(x) < 33 or ord(x) > 126 for x in bech)) or (bech.lower() != bech and bech.upper() != bech)):
65-
return (None, None)
66+
raise ValueError("Not a bech32-encoded string: {}".format(bech))
67+
6668
bech = bech.lower()
6769
pos = bech.rfind('1')
6870
if pos < 1 or pos + 7 > len(bech):
69-
return (None, None)
71+
raise ValueError("Could not locate hrp separator '1' in {}".format(bech))
72+
7073
if not all(x in CHARSET for x in bech[pos + 1:]):
71-
return (None, None)
74+
raise ValueError("Non-bech32 character found in {}".format(bech))
75+
7276
hrp = bech[:pos]
73-
data = [CHARSET.find(x) for x in bech[pos + 1:]]
77+
data = bytes([CHARSET.find(x) for x in bech[pos + 1:]])
7478
if not bech32_verify_checksum(hrp, data):
75-
return (None, None)
79+
raise ValueError("Checksum verification failed for {}".format(bech))
80+
7681
return (hrp, data[:-6])
7782

7883

contrib/pyln-proto/pyln/proto/invoice.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -59,20 +59,20 @@ def unshorten_amount(amount):
5959

6060

6161
# Bech32 spits out array of 5-bit values. Shim here.
62-
def u5_to_bitarray(arr):
62+
def u5_to_bitarray(arr: bytes):
6363
ret = bitstring.BitArray()
6464
for a in arr:
6565
ret += bitstring.pack("uint:5", a)
6666
return ret
6767

6868

69-
def bitarray_to_u5(barr):
69+
def bitarray_to_u5(barr) -> bytes:
7070
assert barr.len % 5 == 0
7171
ret = []
7272
s = bitstring.ConstBitStream(barr)
7373
while s.pos != s.len:
7474
ret.append(s.read(5).uint)
75-
return ret
75+
return bytes(ret)
7676

7777

7878
def encode_fallback(fallback, currency):

contrib/pyln-proto/pyln/proto/zbase32.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import bitstring
1+
import bitstring # type: ignore
22

33

44
zbase32_chars = b'ybndrfg8ejkmcpqxot1uwisza345h769'

0 commit comments

Comments
 (0)