From 9546c7fc716c88c823ef9d9f1f8395f82da70658 Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Tue, 12 Dec 2023 15:36:39 +0200
Subject: [PATCH] improved type annotations

---
 kyber/encryption/decrypt.py | 2 +-
 kyber/encryption/encrypt.py | 2 +-
 kyber/encryption/keygen.py  | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/kyber/encryption/decrypt.py b/kyber/encryption/decrypt.py
index 7453560..78db422 100644
--- 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
index d3e0e71..56bde90 100644
--- 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
index 6a4b488..fb9fa7f 100644
--- 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)