Improve readability of ciphers/mixed_keyword_cypher.py

ciphers/mixed_keyword_cypher.py

@@ -1,7 +1,9 @@
-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, alphabet: str = ascii_uppercase) -> str:
+    """
+    For keyword: hello
 
     H E L O
     A B C D
@@ -19,50 +21,48 @@ def mixed_keyword(key: str = "college", pt: str = "UNIVERSITY") -> str:
      'Y': 'T', 'Z': 'Y'}
     '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
+    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
+
+    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)
 
 
-print(mixed_keyword("college", "UNIVERSITY"))
+if __name__ == "__main__":
+    # example use
+    print(mixed_keyword("college", "UNIVERSITY", ascii_uppercase))