Browse Source

pyln: Adds type annotations to zbase32 functions

travis-experimental
Sergi Delgado Segura 4 years ago
committed by Christian Decker
parent
commit
9de5f438c4
  1. 13
      contrib/pyln-proto/pyln/proto/zbase32.py

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

@ -1,4 +1,5 @@
import bitstring # type: ignore
from typing import Union
zbase32_chars = b'ybndrfg8ejkmcpqxot1uwisza345h769'
@ -23,7 +24,7 @@ zbase32_revchars = [
]
def _message_to_bitarray(message):
def _message_to_bitarray(message: bytes) -> bitstring.ConstBitStream:
"""Encodes a message as a bitarray with length multiple of 5."""
barr = bitstring.ConstBitStream(message)
padding_len = 5 - (len(barr) % 5)
@ -42,7 +43,7 @@ def _bitarray_to_message(barr):
return barr.bytes
def _bitarray_to_u5(barr):
def _bitarray_to_u5(barr: bitstring.ConstBitStream) -> list:
"""Converts a bitarray in a list of uint5."""
ret = []
while barr.pos != barr.len:
@ -50,7 +51,7 @@ def _bitarray_to_u5(barr):
return ret
def _u5_to_bitarray(arr):
def _u5_to_bitarray(arr: list) -> bitstring.BitArray:
"""Converts a list of uint5 values to a bitarray."""
ret = bitstring.BitArray()
for a in arr:
@ -58,7 +59,7 @@ def _u5_to_bitarray(arr):
return ret
def is_zbase32_encoded(message):
def is_zbase32_encoded(message: Union[str, bytes]) -> bool:
"""Checks if a message is zbase32 encoded."""
if isinstance(message, str):
message = message.encode("ASCII")
@ -67,7 +68,7 @@ def is_zbase32_encoded(message):
return set(message).issubset(zbase32_chars)
def encode(message):
def encode(message: Union[str, bytes]) -> bytes:
"""Encodes a message (str or bytes) to zbase32."""
if isinstance(message, str):
message = message.encode('ASCII')
@ -80,7 +81,7 @@ def encode(message):
return bytes(res)
def decode(message):
def decode(message: Union[str, bytes]) -> bytes:
"""Decodes a message (str or bytes) from zbase32."""
if isinstance(message, str):
message = message.encode('ASCII')

Loading…
Cancel
Save