Add tests and type hints to hill_cipher.py

ciphers/hill_cipher.py

@@ -42,6 +42,10 @@
 
 
 def gcd(a, b):
+    """
+    >>> gcd(2, 5)
+    1
+    """
     if a == 0:
         return b
     return gcd(b % a, a)
@@ -52,9 +56,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)
 
@@ -63,14 +64,34 @@ class HillCipher:
     def __init__(self, encrypt_key):
         """
         encrypt_key is an NxN numpy matrix
+        >>> HillCipher(numpy.matrix([[2,5],[1,6]])).__init__(numpy.matrix([[2,5],[1,6]]))
+        >>>
         """
         self.encrypt_key = self.modulus(encrypt_key)  # mod36 calc's on the encrypt key
         self.check_determinant()  # validate the determinant of the encryption key
         self.decrypt_key = None
         self.break_key = encrypt_key.shape[0]
 
+    def replace_letters(self, letter):
+        """
+        >>> HillCipher(numpy.matrix([[2,5],[1,6]])).replace_letters('T')
+        19
+        """
+        return self.key_string.index(letter)
+
+    def replace_numbers(self, num):
+        """
+        >>> HillCipher(numpy.matrix([[2,5],[1,6]])).replace_numbers(19)
+        'T'
+        """
+        return self.key_string[int(round(num))]
+
     def check_determinant(self):
-        det = round(numpy.linalg.det(self.encrypt_key))
+        """
+        >>> HillCipher(numpy.matrix([[2,5],[1,6]])).check_determinant()
+        >>>
+        """
+        det: int = round(numpy.linalg.det(self.encrypt_key))
 
         if det < 0:
             det = det % len(self.key_string)
@@ -84,8 +105,13 @@ def check_determinant(self):
             )
 
     def process_text(self, text):
-        text = list(text.upper())
-        text = [char for char in text if char in self.key_string]
+        """
+        >>> HillCipher(numpy.matrix([[2,5],[1,6]])).process_text('testing hill cipher')
+        'TESTINGHILLCIPHERR'
+        >>>
+        """
+        text: list = list(text.upper())
+        text: list = [char for char in text if char in self.key_string]
 
         last = text[-1]
         while len(text) % self.break_key != 0:
@@ -94,26 +120,37 @@ def process_text(self, text):
         return "".join(text)
 
     def encrypt(self, text):
+        """
+        >>> HillCipher(numpy.matrix([[2,5],[1,6]])).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 = list(map(HillCipher(self.encrypt_key).replace_letters, 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(
+                list(map(HillCipher(self.encrypt_key).replace_numbers, batch_encrypted))
+            )
             encrypted += encrypted_batch
 
         return encrypted
 
     def make_decrypt_key(self):
-        det = round(numpy.linalg.det(self.encrypt_key))
+        """
+        >>> HillCipher(numpy.matrix([[2,5],[1,6]])).make_decrypt_key()
+        matrix([[ 6., 25.],
+                [ 5., 26.]])
+        """
+        det: int = round(numpy.linalg.det(self.encrypt_key))
 
         if det < 0:
-            det = det % len(self.key_string)
+            det: int = det % len(self.key_string)
         det_inv = None
         for i in range(len(self.key_string)):
             if (det * i) % len(self.key_string) == 1:
@@ -129,32 +166,55 @@ def make_decrypt_key(self):
         return self.toInt(self.modulus(inv_key))
 
     def decrypt(self, text):
+        """
+        >>> HillCipher(numpy.matrix([[2,5],[1,6]])).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 = list(map(HillCipher(self.encrypt_key).replace_letters, 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(
+                list(map(HillCipher(self.encrypt_key).replace_numbers, batch_decrypted))
+            )
             decrypted += decrypted_batch
 
         return decrypted
 
 
-def main():
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()
+    """
+    Enter the order of the encryption key: 2
+Enter each row of the encryption key with space separated integers
+'2 5'
+'1 6'
+Would you like to encrypt or decrypt some text? (1 or 2)
+
+1. Encrypt
+2. Decrypt
+1
+What text would you like to encrypt?: 'testing hill cipher'
+Your encrypted text is:
+WHXYJOLM9C6XT085LL
+"""
     N = int(input("Enter the order of the encryption key: "))
-    hill_matrix = []
+    hill_matrix: list = []
 
     print("Enter each row of the encryption key with space separated integers")
     for i in range(N):
         row = list(map(int, input().split()))
         hill_matrix.append(row)
-
+    print(hill_matrix)
     hc = HillCipher(numpy.matrix(hill_matrix))
 
     print("Would you like to encrypt or decrypt some text? (1 or 2)")
@@ -165,15 +225,11 @@ def main():
 """
     )
 
-    if option == "1":
+    if option == 1:
         text_e = input("What text would you like to encrypt?: ")
         print("Your encrypted text is:")
         print(hc.encrypt(text_e))
-    elif option == "2":
+    elif option == 2:
         text_d = input("What text would you like to decrypt?: ")
         print("Your decrypted text is:")
         print(hc.decrypt(text_d))
-
-
-if __name__ == "__main__":
-    main()