BIP-0340: Add typing annotations to reference.py

Passes mypy's strict-mode with mypy 0.770.

Author: ysangkok

bip-0340/reference.py

@@ -1,3 +1,4 @@
+from typing import Tuple, Optional, Any
 import hashlib
 import binascii
 
@@ -17,22 +18,24 @@
 # represented by the None keyword.
 G = (0x79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798, 0x483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8)
 
+Point = Tuple[int, int]
+
 # This implementation can be sped up by storing the midstate after hashing
 # tag_hash instead of rehashing it all the time.
-def tagged_hash(tag, msg):
+def tagged_hash(tag: str, msg: bytes) -> bytes:
     tag_hash = hashlib.sha256(tag.encode()).digest()
     return hashlib.sha256(tag_hash + tag_hash + msg).digest()
 
-def is_infinity(P):
+def is_infinity(P: Optional[Point]) -> bool:
     return P is None
 
-def x(P):
+def x(P: Point) -> int:
     return P[0]
 
-def y(P):
+def y(P: Point) -> int:
     return P[1]
 
-def point_add(P1, P2):
+def point_add(P1: Optional[Point], P2: Optional[Point]) -> Optional[Point]:
     if P1 is None:
         return P2
     if P2 is None:
@@ -46,24 +49,24 @@ def point_add(P1, P2):
     x3 = (lam * lam - x(P1) - x(P2)) % p
     return (x3, (lam * (x(P1) - x3) - y(P1)) % p)
 
-def point_mul(P, n):
+def point_mul(P: Optional[Point], n: int) -> Optional[Point]:
     R = None
     for i in range(256):
         if (n >> i) & 1:
             R = point_add(R, P)
         P = point_add(P, P)
     return R
 
-def bytes_from_int(x):
+def bytes_from_int(x: int) -> bytes:
     return x.to_bytes(32, byteorder="big")
 
-def bytes_from_point(P):
+def bytes_from_point(P: Point) -> bytes:
     return bytes_from_int(x(P))
 
-def xor_bytes(b0, b1):
+def xor_bytes(b0: bytes, b1: bytes) -> bytes:
     return bytes(x ^ y for (x, y) in zip(b0, b1))
 
-def lift_x_square_y(b):
+def lift_x_square_y(b: bytes) -> Optional[Point]:
     x = int_from_bytes(b)
     if x >= p:
         return None
@@ -73,36 +76,40 @@ def lift_x_square_y(b):
         return None
     return (x, y)
 
-def lift_x_even_y(b):
+def lift_x_even_y(b: bytes) -> Optional[Point]:
     P = lift_x_square_y(b)
     if P is None:
         return None
     else:
         return (x(P), y(P) if y(P) % 2 == 0 else p - y(P))
 
-def int_from_bytes(b):
+def int_from_bytes(b: bytes) -> int:
     return int.from_bytes(b, byteorder="big")
 
-def hash_sha256(b):
+def hash_sha256(b: bytes) -> bytes:
     return hashlib.sha256(b).digest()
 
-def is_square(x):
-    return pow(x, (p - 1) // 2, p) == 1
+def is_square(x: int) -> bool:
+    return int(pow(x, (p - 1) // 2, p)) == 1
 
-def has_square_y(P):
-    return (not is_infinity(P)) and is_square(y(P))
+def has_square_y(P: Optional[Point]) -> bool:
+    infinity = is_infinity(P)
+    if infinity: return False
+    assert P is not None
+    return is_square(y(P))
 
-def has_even_y(P):
+def has_even_y(P: Point) -> bool:
     return y(P) % 2 == 0
 
-def pubkey_gen(seckey):
+def pubkey_gen(seckey: bytes) -> bytes:
     d0 = int_from_bytes(seckey)
     if not (1 <= d0 <= n - 1):
         raise ValueError('The secret key must be an integer in the range 1..n-1.')
     P = point_mul(G, d0)
+    assert P is not None
     return bytes_from_point(P)
 
-def schnorr_sign(msg, seckey, aux_rand):
+def schnorr_sign(msg: bytes, seckey: bytes, aux_rand: bytes) -> bytes:
     if len(msg) != 32:
         raise ValueError('The message must be a 32-byte array.')
     d0 = int_from_bytes(seckey)
@@ -111,12 +118,14 @@ def schnorr_sign(msg, seckey, aux_rand):
     if len(aux_rand) != 32:
         raise ValueError('aux_rand must be 32 bytes instead of %i.' % len(aux_rand))
     P = point_mul(G, d0)
+    assert P is not None
     d = d0 if has_even_y(P) else n - d0
     t = xor_bytes(bytes_from_int(d), tagged_hash("BIP340/aux", aux_rand))
     k0 = int_from_bytes(tagged_hash("BIP340/nonce", t + bytes_from_point(P) + msg)) % n
     if k0 == 0:
         raise RuntimeError('Failure. This happens only with negligible probability.')
     R = point_mul(G, k0)
+    assert R is not None
     k = n - k0 if not has_square_y(R) else k0
     e = int_from_bytes(tagged_hash("BIP340/challenge", bytes_from_point(R) + bytes_from_point(P) + msg)) % n
     sig = bytes_from_point(R) + bytes_from_int((k + e * d) % n)
@@ -125,7 +134,7 @@ def schnorr_sign(msg, seckey, aux_rand):
         raise RuntimeError('The created signature does not pass verification.')
     return sig
 
-def schnorr_verify(msg, pubkey, sig):
+def schnorr_verify(msg: bytes, pubkey: bytes, sig: bytes) -> bool:
     if len(msg) != 32:
         raise ValueError('The message must be a 32-byte array.')
     if len(pubkey) != 32:
@@ -153,26 +162,26 @@ def schnorr_verify(msg, pubkey, sig):
 import os
 import sys
 
-def test_vectors():
+def test_vectors() -> bool:
     all_passed = True
     with open(os.path.join(sys.path[0], 'test-vectors.csv'), newline='') as csvfile:
         reader = csv.reader(csvfile)
         reader.__next__()
         for row in reader:
-            (index, seckey, pubkey, aux_rand, msg, sig, result, comment) = row
-            pubkey = bytes.fromhex(pubkey)
-            msg = bytes.fromhex(msg)
-            sig = bytes.fromhex(sig)
-            result = result == 'TRUE'
+            (index, seckey_hex, pubkey_hex, aux_rand_hex, msg_hex, sig_hex, result_str, comment) = row
+            pubkey = bytes.fromhex(pubkey_hex)
+            msg = bytes.fromhex(msg_hex)
+            sig = bytes.fromhex(sig_hex)
+            result = result_str == 'TRUE'
             print('\nTest vector', ('#' + index).rjust(3, ' ') + ':')
-            if seckey != '':
-                seckey = bytes.fromhex(seckey)
+            if seckey_hex != '':
+                seckey = bytes.fromhex(seckey_hex)
                 pubkey_actual = pubkey_gen(seckey)
                 if pubkey != pubkey_actual:
                     print(' * Failed key generation.')
                     print('   Expected key:', pubkey.hex().upper())
                     print('     Actual key:', pubkey_actual.hex().upper())
-                aux_rand = bytes.fromhex(aux_rand)
+                aux_rand = bytes.fromhex(aux_rand_hex)
                 try:
                     sig_actual = schnorr_sign(msg, seckey, aux_rand)
                     if sig == sig_actual:
@@ -207,7 +216,7 @@ def test_vectors():
 #
 import inspect
 
-def pretty(v):
+def pretty(v: Any) -> Any:
     if isinstance(v, bytes):
         return '0x' + v.hex()
     if isinstance(v, int):
@@ -216,9 +225,12 @@ def pretty(v):
         return tuple(map(pretty, v))
     return v
 
-def debug_print_vars():
+def debug_print_vars() -> None:
     if DEBUG:
-        frame = inspect.currentframe().f_back
+        current_frame = inspect.currentframe()
+        assert current_frame is not None
+        frame = current_frame.f_back
+        assert frame is not None
         print('   Variables in function ', frame.f_code.co_name, ' at line ', frame.f_lineno, ':', sep='')
         for var_name, var_val in frame.f_locals.items():
             print('   ' + var_name.rjust(11, ' '), '==', pretty(var_val))