Add docstrings and doctests and fix a bug ciphers/trifid_cipher.py

ciphers/trifid_cipher.py

@@ -1,8 +1,28 @@
-# https://en.wikipedia.org/wiki/Trifid_cipher
+"""
+The trifid cipher uses a table to fractionate each plaintext letter into a
+trigram,mixes the constituents of the trigrams, and then applies the table
+in reverse to turn these mixed trigrams into ciphertext letters.
+https://en.wikipedia.org/wiki/Trifid_cipher
+"""
+
 from __future__ import annotations
 
 
 def __encrypt_part(message_part: str, character_to_number: dict[str, str]) -> str:
+    """
+    Arranges the triagram value of each letter of 'message_part' vertically
+    and joins them horizontally
+
+    >>> __encrypt_part('ASK',
+    ... {'A': '111', 'B': '112', 'C':'113', 'D': '121', 'E': '122', 'F': '123',
+    ... 'G': '131', 'H': '132', 'I': '133',
+    ... 'J': '211', 'K': '212', 'L': '213', 'M': '221', 'N': '222',
+    ... 'O': '223', 'P': '231', 'Q': '232', 'R': '233', 'S': '311', 'T': '312',
+    ... 'U': '313', 'V': '321', 'W': '322', 'X': '323', 'Y': '331', 'Z': '332',
+    ... '+': '333'})
+    '132111112'
+
+    """
     one, two, three = "", "", ""
     tmp = []
 
@@ -20,6 +40,20 @@ def __encrypt_part(message_part: str, character_to_number: dict[str, str]) -> st
 def __decrypt_part(
     message_part: str, character_to_number: dict[str, str]
 ) -> tuple[str, str, str]:
+    """
+    Converts each letter of the input string into there respective trigram
+    values, joins them and splits them into three equal groups of strings.
+    Then returns the group of strings .
+
+    >>> __decrypt_part('ABCDE',
+    ... {'A': '111', 'B': '112', 'C':'113', 'D': '121', 'E': '122', 'F': '123',
+    ... 'G': '131', 'H': '132', 'I': '133',
+    ... 'J': '211', 'K': '212', 'L': '213', 'M': '221', 'N': '222',
+    ... 'O': '223', 'P': '231', 'Q': '232', 'R': '233', 'S': '311', 'T': '312',
+    ... 'U': '313', 'V': '321', 'W': '322', 'X': '323', 'Y': '331', 'Z': '332',
+    ... '+': '333'})
+    ('11111', '21131', '21122')
+    """
     tmp, this_part = "", ""
     result = []
 
@@ -38,6 +72,52 @@ def __decrypt_part(
 def __prepare(
     message: str, alphabet: str
 ) -> tuple[str, str, dict[str, str], dict[str, str]]:
+    """
+    A helper function that generates the triagrams and assigns each letter
+    of the alphabet to its corresponding triagram and stores this in a
+    dictionary ("character_to_number" and "number_to_character") after
+    confirming if the alphabet's length is 27.
+
+    >>> test = __prepare('I aM a BOy','abCdeFghijkLmnopqrStuVwxYZ+')
+    >>> expected = ('IAMABOY','ABCDEFGHIJKLMNOPQRSTUVWXYZ+',
+    ... {'A': '111', 'B': '112', 'C':'113', 'D': '121', 'E': '122', 'F': '123',
+    ... 'G': '131', 'H': '132', 'I': '133',
+    ... 'J': '211', 'K': '212', 'L': '213', 'M': '221', 'N': '222',
+    ... 'O': '223', 'P': '231', 'Q': '232', 'R': '233', 'S': '311', 'T': '312',
+    ... 'U': '313', 'V': '321', 'W': '322', 'X': '323', 'Y': '331', 'Z': '332',
+    ... '+': '333'},
+    ... {'111': 'A', '112': 'B', '113': 'C', '121': 'D', '122': 'E',
+    ... '123': 'F', '131': 'G', '132': 'H', '133': 'I', '211': 'J', '212': 'K',
+    ... '213': 'L', '221': 'M', '222': 'N', '223': 'O', '231': 'P', '232': 'Q',
+    ... '233': 'R', '311': 'S', '312': 'T', '313': 'U', '321': 'V', '322': 'W',
+    ... '323': 'X', '331': 'Y', '332': 'Z', '333': '+'})
+    >>> test == expected
+    True
+
+    Testing with incomplete alphabet
+    >>> __prepare('I aM a BOy','abCdeFghijkLmnopqrStuVw')
+    Traceback (most recent call last):
+        ...
+    KeyError: 'Length of alphabet has to be 27.'
+
+    Testing with extra long alphabets
+    >>> __prepare('I aM a BOy','abCdeFghijkLmnopqrStuVwxyzzwwtyyujjgfd')
+    Traceback (most recent call last):
+        ...
+    KeyError: 'Length of alphabet has to be 27.'
+
+    Testing with punctuations that are not in the given alphabet
+    >>> __prepare('am i a boy?','abCdeFghijkLmnopqrStuVwxYZ+')
+    Traceback (most recent call last):
+        ...
+    ValueError: Each message character has to be included in alphabet!
+
+    Testing with numbers
+    >>> __prepare(500,'abCdeFghijkLmnopqrStuVwxYZ+')
+    Traceback (most recent call last):
+        ...
+    AttributeError: 'int' object has no attribute 'replace'
+    """
     # Validate message and alphabet, set to upper and remove spaces
     alphabet = alphabet.replace(" ", "").upper()
     message = message.replace(" ", "").upper()
@@ -91,6 +171,32 @@ def __prepare(
 def encrypt_message(
     message: str, alphabet: str = "ABCDEFGHIJKLMNOPQRSTUVWXYZ.", period: int = 5
 ) -> str:
+    """
+    encrypt_message
+    ===============
+
+    Encrypts a message using the trifid_cipher. Any punctuatuions that
+    would be used should be added to the alphabet.
+
+    PARAMETERS
+    ----------
+
+    *   message: The message you want to encrypt.
+    *   alphabet (optional): The characters to be used for the cipher .
+    *   period (optional): The number of characters you want in a group whilst
+        encrypting.
+
+    >>> encrypt_message('I am a boy')
+    'BCDGBQY'
+
+    >>> encrypt_message(' ')
+    ''
+
+    >>> encrypt_message('   aide toi le c  iel      ta id  era    ',
+    ... 'FELIXMARDSTBCGHJKNOPQUVWYZ+',5)
+    'FMJFVOISSUFTFPUFEQQC'
+
+    """
     message, alphabet, character_to_number, number_to_character = __prepare(
         message, alphabet
     )
@@ -110,13 +216,35 @@ def encrypt_message(
 def decrypt_message(
     message: str, alphabet: str = "ABCDEFGHIJKLMNOPQRSTUVWXYZ.", period: int = 5
 ) -> str:
+    """
+    decrypt_message
+    ===============
+
+    Decrypts a trifid_cipher encrypted message .
+
+    PARAMETERS
+    ----------
+
+    *   message: The message you want to decrypt .
+    *   alphabet (optional): The characters used for the cipher.
+    *   period (optional): The number of characters used in grouping when it
+        was encrypted.
+
+    >>> decrypt_message('BCDGBQY')
+    'IAMABOY'
+
+    Decrypting with your own alphabet and period
+    >>> decrypt_message('FMJFVOISSUFTFPUFEQQC','FELIXMARDSTBCGHJKNOPQUVWYZ+',5)
+    'AIDETOILECIELTAIDERA'
+    """
     message, alphabet, character_to_number, number_to_character = __prepare(
         message, alphabet
     )
+
     decrypted_numeric = []
     decrypted = ""
 
-    for i in range(0, len(message) + 1, period):
+    for i in range(0, len(message), period):
         a, b, c = __decrypt_part(message[i : i + period], character_to_number)
 
         for j in range(len(a)):
@@ -129,6 +257,9 @@ def decrypt_message(
 
 
 if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()
     msg = "DEFEND THE EAST WALL OF THE CASTLE."
     encrypted = encrypt_message(msg, "EPSDUCVWYM.ZLKXNBTFGORIJHAQ")
     decrypted = decrypt_message(encrypted, "EPSDUCVWYM.ZLKXNBTFGORIJHAQ")