From 9de5f438c4e4c09dd07e179dc8325d06183253cf Mon Sep 17 00:00:00 2001
From: Sergi Delgado Segura <sergi.delgado.s@gmail.com>
Date: Wed, 21 Oct 2020 18:15:36 +0200
Subject: [PATCH] pyln: Adds type annotations to zbase32 functions

---
 contrib/pyln-proto/pyln/proto/zbase32.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/contrib/pyln-proto/pyln/proto/zbase32.py b/contrib/pyln-proto/pyln/proto/zbase32.py
index 75b017d35..22a575a15 100644
--- a/contrib/pyln-proto/pyln/proto/zbase32.py
+++ b/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')