improved type annotations

PyryL · Dec 12, 2023 · 9546c7f · 9546c7f
1 parent 4565eed
commit 9546c7f
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/kyber/encryption/decrypt.py b/kyber/encryption/decrypt.py
@@ -4,7 +4,7 @@
 from kyber.constants import n, k, du, dv
 from kyber.entities.polring import PolynomialRing
 
-def decrypt(private_key, ciphertext) -> bytes:
+def decrypt(private_key: bytes, ciphertext: bytes) -> bytes:
     """
     Decrypts the given ciphertext with the given private key.
     :returns Decrypted 32-bit shared secret

diff --git a/kyber/encryption/encrypt.py b/kyber/encryption/encrypt.py
@@ -25,7 +25,7 @@ def secret(self) -> bytes:
         """The 32-bit shared secret that was encrypted."""
         return self._m
 
-    def encrypt(self):
+    def encrypt(self) -> bytes:
         """
         Encrypts 32-bit random shared secret.
         :returns Ciphertext

diff --git a/kyber/encryption/keygen.py b/kyber/encryption/keygen.py
@@ -8,7 +8,7 @@
 from kyber.utils.parse import parse
 from kyber.entities.polring import PolynomialRing
 
-def generate_keys() -> tuple:
+def generate_keys() -> tuple[bytes, bytes]:
     """
     Generates a new Kyber keypair.
     :returns (private_key, public_key)