Improve readability of ciphers/mixed_keyword_cypher.py

ciphers/mixed_keyword_cypher.py

@@ -1,7 +1,11 @@
-def mixed_keyword(key: str = "college", pt: str = "UNIVERSITY") -> str:
-    """
+from string import ascii_uppercase
+
 
-    For key:hello
+def mixed_keyword(
+    keyword: str, plaintext: str, verbose: bool = True, alphabet: str = ascii_uppercase
+) -> str:
+    """
+    For keyword: hello
 
     H E L O
     A B C D
@@ -12,58 +16,60 @@ def mixed_keyword(key: str = "college", pt: str = "UNIVERSITY") -> str:
     Y Z
     and map vertically
 
-    >>> mixed_keyword("college", "UNIVERSITY")  # doctest: +NORMALIZE_WHITESPACE
+    >>> mixed_keyword("college", "UNIVERSITY", True)  # doctest: +NORMALIZE_WHITESPACE
     {'A': 'C', 'B': 'A', 'C': 'I', 'D': 'P', 'E': 'U', 'F': 'Z', 'G': 'O', 'H': 'B',
      'I': 'J', 'J': 'Q', 'K': 'V', 'L': 'L', 'M': 'D', 'N': 'K', 'O': 'R', 'P': 'W',
      'Q': 'E', 'R': 'F', 'S': 'M', 'T': 'S', 'U': 'X', 'V': 'G', 'W': 'H', 'X': 'N',
      'Y': 'T', 'Z': 'Y'}
     'XKJGUFMJST'
+
+    >>> mixed_keyword("college", "UNIVERSITY", False)  # doctest: +NORMALIZE_WHITESPACE
+    'XKJGUFMJST'
     """
-    key = key.upper()
-    pt = pt.upper()
-    temp = []
-    for i in key:
-        if i not in temp:
-            temp.append(i)
-    len_temp = len(temp)
-    # print(temp)
-    alpha = []
-    modalpha = []
-    for j in range(65, 91):
-        t = chr(j)
-        alpha.append(t)
-        if t not in temp:
-            temp.append(t)
-    # print(temp)
-    r = int(26 / 4)
-    # print(r)
-    k = 0
-    for _ in range(r):
-        s = []
-        for _ in range(len_temp):
-            s.append(temp[k])
-            if k >= 25:
-                break
-            k += 1
-        modalpha.append(s)
-    # print(modalpha)
-    d = {}
-    j = 0
-    k = 0
-    for j in range(len_temp):
-        for m in modalpha:
-            if not len(m) - 1 >= j:
-                break
-            d[alpha[k]] = m[j]
-            if not k < 25:
+    keyword = keyword.upper()
+    plaintext = plaintext.upper()
+    alphabet_set = set(alphabet)
+
+    # create a list of unique characters in the keyword - their order matters
+    # it determines how we will map plaintext characters to the ciphertext
+    unique_chars = []
+    for char in keyword:
+        if char in alphabet_set and char not in unique_chars:
+            unique_chars.append(char)
+    # the number of those unique characters will determine the number of rows
+    num_unique_chars_in_keyword = len(unique_chars)
+
+    # create a shifted version of the alphabet
+    shifted_alphabet = unique_chars + [
+        char for char in alphabet if char not in unique_chars
+    ]
+
+    # create a modified alphabet by splitting the shifted alphabet into rows
+    modified_alphabet = [
+        shifted_alphabet[k : k + num_unique_chars_in_keyword]
+        for k in range(0, 26, num_unique_chars_in_keyword)
+    ]
+
+    # map the alphabet characters to the modified alphabet characters
+    # going 'vertically' through the modified alphabet - consider columns first
+    mapping = {}
+    letter_index = 0
+    for column in range(num_unique_chars_in_keyword):
+        for row in modified_alphabet:
+            # if current row (the last one) is too short, break out of loop
+            if len(row) <= column:
                 break
-            k += 1
-    print(d)
-    cypher = ""
-    for i in pt:
-        cypher += d[i]
-    return cypher
+
+            # map current letter to letter in modified alphabet
+            mapping[alphabet[letter_index]] = row[column]
+            letter_index += 1
+
+    if verbose:
+        print(mapping)
+    # create the encrypted text by mapping the plaintext to the modified alphabet
+    return "".join(mapping[char] if char in mapping else char for char in plaintext)
 
 
 if __name__ == "__main__":
+    # example use
     print(mixed_keyword("college", "UNIVERSITY"))