Add aliquot_sum.py new function and doctests

ciphers/atbash.py

@@ -1,54 +1,75 @@
 """https://en.wikipedia.org/wiki/Atbash"""
 
 import string
+from timeit import timeit  # Moved to top-level as required
 
 
 def atbash_slow(sequence: str) -> str:
     """
+    Atbash cipher implementation using ordinal values.
+    Encodes/decodes by reversing the alphabet.
+
     >>> atbash_slow("ABCDEFG")
     'ZYXWVUT'
 
     >>> atbash_slow("aW;;123BX")
     'zD;;123YC'
     """
     output = ""
-    for i in sequence:
-        extract = ord(i)
-        if 65 <= extract <= 90:
-            output += chr(155 - extract)
-        elif 97 <= extract <= 122:
-            output += chr(219 - extract)
-        else:
-            output += i
+    for char in sequence:
+        code = ord(char)
+        if 65 <= code <= 90:  # Uppercase A-Z
+            output += chr(155 - code)
+        elif 97 <= code <= 122:  # Lowercase a-z
+            output += chr(219 - code)
+        else:  # Non-alphabetic characters
+            output += char
     return output
 
 
 def atbash(sequence: str) -> str:
     """
+    Optimized Atbash cipher implementation using string translation.
+    More efficient than ordinal-based approach.
+
     >>> atbash("ABCDEFG")
     'ZYXWVUT'
 
     >>> atbash("aW;;123BX")
     'zD;;123YC'
     """
+    # Create translation tables
     letters = string.ascii_letters
-    letters_reversed = string.ascii_lowercase[::-1] + string.ascii_uppercase[::-1]
-    return "".join(
-        letters_reversed[letters.index(c)] if c in letters else c for c in sequence
-    )
+    reversed_letters = string.ascii_lowercase[::-1] + string.ascii_uppercase[::-1]
 
+    # Create translation mapping
+    translation = str.maketrans(letters, reversed_letters)
 
-def benchmark() -> None:
-    """Let's benchmark our functions side-by-side..."""
-    from timeit import timeit
+    # Apply translation to each character
+    return sequence.translate(translation)
 
+
+def benchmark() -> None:
+    """
+    Performance comparison of both Atbash implementations.
+    Measures execution time using Python's timeit module.
+    """
     print("Running performance benchmarks...")
-    setup = "from string import printable ; from __main__ import atbash, atbash_slow"
-    print(f"> atbash_slow(): {timeit('atbash_slow(printable)', setup=setup)} seconds")
-    print(f">      atbash(): {timeit('atbash(printable)', setup=setup)} seconds")
+    setup = "from string import printable; from __main__ import atbash, atbash_slow"
+    # Time the slow implementation
+    slow_time = timeit("atbash_slow(printable)", setup=setup)
+    print(f"> atbash_slow(): {slow_time:.6f} seconds")
+
+    # Time the optimized implementation
+    fast_time = timeit("atbash(printable)", setup=setup)
+    print(f">      atbash(): {fast_time:.6f} seconds")
 
 
 if __name__ == "__main__":
-    for example in ("ABCDEFGH", "123GGjj", "testStringtest", "with space"):
+    # Test examples
+    examples = ("ABCDEFGH", "123GGjj", "testStringtest", "with space")
+    for example in examples:
         print(f"{example} encrypted in atbash: {atbash(example)}")
+
+    # Run performance comparison
     benchmark()

ciphers/hill_cipher.py

@@ -44,176 +44,111 @@
 
 
 class HillCipher:
-    key_string = string.ascii_uppercase + string.digits
-    # This cipher takes alphanumerics into account
-    # i.e. a total of 36 characters
-
-    # take x and return x % len(key_string)
-    modulus = np.vectorize(lambda x: x % 36)
-
-    to_int = np.vectorize(round)
+    key_string = string.ascii_uppercase + string.digits  # 36 chars
+    modulus = np.vectorize(lambda x: x % 36)  # Mod 36
+    to_int = np.vectorize(round)  # Round numbers
 
     def __init__(self, encrypt_key: np.ndarray) -> None:
-        """
-        encrypt_key is an NxN numpy array
-        """
-        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.break_key = encrypt_key.shape[0]
+        self.encrypt_key = self.modulus(encrypt_key)
+        self.check_determinant()  # Validate key
+        self.break_key = encrypt_key.shape[0]  # Matrix size
 
     def replace_letters(self, letter: str) -> int:
-        """
-        >>> hill_cipher = HillCipher(np.array([[2, 5], [1, 6]]))
-        >>> hill_cipher.replace_letters('T')
-        19
-        >>> hill_cipher.replace_letters('0')
-        26
-        """
+        """Char to index (A=0, 0=26)"""
         return self.key_string.index(letter)
 
     def replace_digits(self, num: int) -> str:
-        """
-        >>> hill_cipher = HillCipher(np.array([[2, 5], [1, 6]]))
-        >>> hill_cipher.replace_digits(19)
-        'T'
-        >>> hill_cipher.replace_digits(26)
-        '0'
-        """
-        return self.key_string[round(num)]
+        """Index to char"""
+        return self.key_string[num]
 
     def check_determinant(self) -> None:
-        """
-        >>> hill_cipher = HillCipher(np.array([[2, 5], [1, 6]]))
-        >>> hill_cipher.check_determinant()
-        """
+        """Ensure det(key) coprime with 36"""
         det = round(np.linalg.det(self.encrypt_key))
-
         if det < 0:
-            det = det % len(self.key_string)
+            det %= len(self.key_string)
 
-        req_l = len(self.key_string)
+        error_msg = f"Det {det} not coprime with 36. Try another key."
         if greatest_common_divisor(det, len(self.key_string)) != 1:
-            msg = (
-                f"determinant modular {req_l} of encryption key({det}) "
-                f"is not co prime w.r.t {req_l}.\nTry another key."
-            )
-            raise ValueError(msg)
+            raise ValueError(error_msg)
 
     def process_text(self, text: str) -> str:
-        """
-        >>> hill_cipher = HillCipher(np.array([[2, 5], [1, 6]]))
-        >>> hill_cipher.process_text('Testing Hill Cipher')
-        'TESTINGHILLCIPHERR'
-        >>> hill_cipher.process_text('hello')
-        'HELLOO'
-        """
-        chars = [char for char in text.upper() if char in self.key_string]
-
-        last = chars[-1]
+        """Uppercase, filter, pad text"""
+        chars = [c for c in text.upper() if c in self.key_string]
+        last = chars[-1] if chars else "A"
         while len(chars) % self.break_key != 0:
             chars.append(last)
-
         return "".join(chars)
 
     def encrypt(self, text: str) -> str:
-        """
-        >>> hill_cipher = HillCipher(np.array([[2, 5], [1, 6]]))
-        >>> hill_cipher.encrypt('testing hill cipher')
-        'WHXYJOLM9C6XT085LL'
-        >>> hill_cipher.encrypt('hello')
-        '85FF00'
-        """
+        """Encrypt with Hill cipher"""
         text = self.process_text(text.upper())
         encrypted = ""
-
-        for i in range(0, len(text) - self.break_key + 1, self.break_key):
+        for i in range(0, len(text), self.break_key):
             batch = text[i : i + self.break_key]
-            vec = [self.replace_letters(char) for char in batch]
+            vec = [self.replace_letters(c) for c in batch]
             batch_vec = np.array([vec]).T
-            batch_encrypted = self.modulus(self.encrypt_key.dot(batch_vec)).T.tolist()[
-                0
-            ]
+            product = self.encrypt_key.dot(batch_vec)
+            modulated = self.modulus(product)
+            batch_encrypted = modulated.T.tolist()[0]
             encrypted_batch = "".join(
-                self.replace_digits(num) for num in batch_encrypted
+                self.replace_digits(round(n)) for n in batch_encrypted
             )
             encrypted += encrypted_batch
-
         return encrypted
 
     def make_decrypt_key(self) -> np.ndarray:
-        """
-        >>> hill_cipher = HillCipher(np.array([[2, 5], [1, 6]]))
-        >>> hill_cipher.make_decrypt_key()
-        array([[ 6, 25],
-               [ 5, 26]])
-        """
+        """Create decryption key"""
         det = round(np.linalg.det(self.encrypt_key))
-
         if det < 0:
-            det = det % len(self.key_string)
-        det_inv = None
-        for i in range(len(self.key_string)):
-            if (det * i) % len(self.key_string) == 1:
-                det_inv = i
-                break
+            det %= len(self.key_string)
+
+        # Find det modular inverse
+        det_inv = next(i for i in range(36) if (det * i) % 36 == 1)
 
+        # Compute inverse key
         inv_key = (
             det_inv * np.linalg.det(self.encrypt_key) * np.linalg.inv(self.encrypt_key)
         )
-
         return self.to_int(self.modulus(inv_key))
 
     def decrypt(self, text: str) -> str:
-        """
-        >>> hill_cipher = HillCipher(np.array([[2, 5], [1, 6]]))
-        >>> hill_cipher.decrypt('WHXYJOLM9C6XT085LL')
-        'TESTINGHILLCIPHERR'
-        >>> hill_cipher.decrypt('85FF00')
-        'HELLOO'
-        """
+        """Decrypt with Hill cipher"""
         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):
+        for i in range(0, len(text), self.break_key):
             batch = text[i : i + self.break_key]
-            vec = [self.replace_letters(char) for char in batch]
+            vec = [self.replace_letters(c) for c in batch]
             batch_vec = np.array([vec]).T
-            batch_decrypted = self.modulus(decrypt_key.dot(batch_vec)).T.tolist()[0]
+            product = decrypt_key.dot(batch_vec)
+            modulated = self.modulus(product)
+            batch_decrypted = modulated.T.tolist()[0]
             decrypted_batch = "".join(
-                self.replace_digits(num) for num in batch_decrypted
+                self.replace_digits(round(n)) for n in batch_decrypted
             )
             decrypted += decrypted_batch
-
         return decrypted
 
 
 def main() -> None:
-    n = int(input("Enter the order of the encryption key: "))
-    hill_matrix = []
+    """Hill Cipher CLI"""
+    n = int(input("Enter key order: "))
+    print(f"Enter {n} rows of space-separated integers:")
+    matrix = [list(map(int, input().split())) for _ in range(n)]
 
-    print("Enter each row of the encryption key with space separated integers")
-    for _ in range(n):
-        row = [int(x) for x in input().split()]
-        hill_matrix.append(row)
+    hc = HillCipher(np.array(matrix))
 
-    hc = HillCipher(np.array(hill_matrix))
+    option = input("1. Encrypt\n2. Decrypt\nChoose: ")
+    text = input("Enter text: ")
 
-    print("Would you like to encrypt or decrypt some text? (1 or 2)")
-    option = input("\n1. Encrypt\n2. Decrypt\n")
     if option == "1":
-        text_e = input("What text would you like to encrypt?: ")
-        print("Your encrypted text is:")
-        print(hc.encrypt(text_e))
+        print("Encrypted:", hc.encrypt(text))
     elif option == "2":
-        text_d = input("What text would you like to decrypt?: ")
-        print("Your decrypted text is:")
-        print(hc.decrypt(text_d))
+        print("Decrypted:", hc.decrypt(text))
 
 
 if __name__ == "__main__":
     import doctest
 
     doctest.testmod()
-
     main()