Add tests and type hints to hill cipher

ciphers/hill_cipher.py

@@ -41,7 +41,17 @@
 import numpy
 
 
-def gcd(a, b):
+def gcd(a: int, b: int) -> int:
+    """
+    >>> gcd(4, 8)
+    4
+    >>> gcd(8, 4)
+    4
+    >>> gcd(4, 7)
+    1
+    >>> gcd(0, 10)
+    10
+    """
     if a == 0:
         return b
     return gcd(b % a, a)
@@ -52,9 +62,6 @@ class HillCipher:
     # This cipher takes alphanumerics into account
     # i.e. a total of 36 characters
 
-    replaceLetters = lambda self, letter: self.key_string.index(letter)
-    replaceNumbers = lambda self, num: self.key_string[round(num)]
-
     # take x and return x % len(key_string)
     modulus = numpy.vectorize(lambda x: x % 36)
 
@@ -69,7 +76,31 @@ def __init__(self, encrypt_key):
         self.decrypt_key = None
         self.break_key = encrypt_key.shape[0]
 
-    def check_determinant(self):
+    def replaceLetters(self, letter: str) -> int:
+        """
+        >>> hill_cipher = HillCipher(numpy.matrix([[2, 5], [1, 6]]))
+        >>> hill_cipher.replaceLetters('T')
+        19
+        >>> hill_cipher.replaceLetters('0')
+        26
+        """
+        return self.key_string.index(letter)
+
+    def replaceNumbers(self, num: int) -> str:
+        """
+        >>> hill_cipher = HillCipher(numpy.matrix([[2, 5], [1, 6]]))
+        >>> hill_cipher.replaceNumbers(19)
+        'T'
+        >>> hill_cipher.replaceNumbers(26)
+        '0'
+        """
+        return self.key_string[round(num)]
+
+    def check_determinant(self) -> None:
+        """
+        >>> hill_cipher = HillCipher(numpy.matrix([[2, 5], [1, 6]]))
+        >>> hill_cipher.check_determinant()
+        """
         det = round(numpy.linalg.det(self.encrypt_key))
 
         if det < 0:
@@ -78,38 +109,54 @@ def check_determinant(self):
         req_l = len(self.key_string)
         if gcd(det, len(self.key_string)) != 1:
             raise ValueError(
-                "discriminant modular {} of encryption key({}) is not co prime w.r.t {}.\nTry another key.".format(
-                    req_l, det, req_l
-                )
+                f"determinant modular {req_l} of encryption key({det}) is not co prime w.r.t {req_l}.\nTry another key."
             )
 
-    def process_text(self, text):
+    def process_text(self, text: str) -> str:
+        """
+        >>> hill_cipher = HillCipher(numpy.matrix([[2, 5], [1, 6]]))
+        >>> hill_cipher.process_text('Testing Hill Cipher')
+        'TESTINGHILLCIPHERR'
+        """
         text = list(text.upper())
-        text = [char for char in text if char in self.key_string]
+        chars = [char for char in text if char in self.key_string]
 
-        last = text[-1]
-        while len(text) % self.break_key != 0:
-            text.append(last)
+        last = chars[-1]
+        while len(chars) % self.break_key != 0:
+            chars.append(last)
 
-        return "".join(text)
+        return "".join(chars)
 
-    def encrypt(self, text):
+    def encrypt(self, text: str) -> str:
+        """
+        >>> hill_cipher = HillCipher(numpy.matrix([[2, 5], [1, 6]]))
+        >>> hill_cipher.encrypt('testing hill cipher')
+        'WHXYJOLM9C6XT085LL'
+        """
         text = self.process_text(text.upper())
         encrypted = ""
 
         for i in range(0, len(text) - self.break_key + 1, self.break_key):
             batch = text[i : i + self.break_key]
-            batch_vec = list(map(self.replaceLetters, batch))
+            batch_vec = [self.replaceLetters(char) for char in batch]
             batch_vec = numpy.matrix([batch_vec]).T
             batch_encrypted = self.modulus(self.encrypt_key.dot(batch_vec)).T.tolist()[
                 0
             ]
-            encrypted_batch = "".join(list(map(self.replaceNumbers, batch_encrypted)))
+            encrypted_batch = "".join(
+                self.replaceNumbers(num) for num in batch_encrypted
+            )
             encrypted += encrypted_batch
 
         return encrypted
 
     def make_decrypt_key(self):
+        """
+        >>> hill_cipher = HillCipher(numpy.matrix([[2, 5], [1, 6]]))
+        >>> hill_cipher.make_decrypt_key()
+        matrix([[ 6., 25.],
+                [ 5., 26.]])
+        """
         det = round(numpy.linalg.det(self.encrypt_key))
 
         if det < 0:
@@ -128,19 +175,26 @@ def make_decrypt_key(self):
 
         return self.toInt(self.modulus(inv_key))
 
-    def decrypt(self, text):
+    def decrypt(self, text: str) -> str:
+        """
+        >>> hill_cipher = HillCipher(numpy.matrix([[2, 5], [1, 6]]))
+        >>> hill_cipher.decrypt('WHXYJOLM9C6XT085LL')
+        'TESTINGHILLCIPHERR'
+        """
         self.decrypt_key = self.make_decrypt_key()
         text = self.process_text(text.upper())
         decrypted = ""
 
         for i in range(0, len(text) - self.break_key + 1, self.break_key):
             batch = text[i : i + self.break_key]
-            batch_vec = list(map(self.replaceLetters, batch))
+            batch_vec = [self.replaceLetters(char) for char in batch]
             batch_vec = numpy.matrix([batch_vec]).T
             batch_decrypted = self.modulus(self.decrypt_key.dot(batch_vec)).T.tolist()[
                 0
             ]
-            decrypted_batch = "".join(list(map(self.replaceNumbers, batch_decrypted)))
+            decrypted_batch = "".join(
+                self.replaceNumbers(num) for num in batch_decrypted
+            )
             decrypted += decrypted_batch
 
         return decrypted
@@ -176,4 +230,7 @@ def main():
 
 
 if __name__ == "__main__":
+    import doctest
+    doctest.testmod()
+
     main()