Baconian cipher with any 2 symbols

Updated to allow  encoding and decoding with any 2 characters

Author: avyukth-fll

ciphers/baconian_cipher.py

@@ -1,89 +1,162 @@
 """
-Program to encode and decode Baconian or Bacon's Cipher
-Wikipedia reference : https://en.wikipedia.org/wiki/Bacon%27s_cipher
+Bacon's/Baconian Cipher Encoder and Decoder
+
+ provides functions to encode and decode messages
+using Bacon's/Baconian cipher with any 2 symbols, not just A & B.
+
+Source: https://en.wikipedia.org/wiki/Bacon's_cipher
 """
 
-encode_dict = {
-    "a": "AAAAA",
-    "b": "AAAAB",
-    "c": "AAABA",
-    "d": "AAABB",
-    "e": "AABAA",
-    "f": "AABAB",
-    "g": "AABBA",
-    "h": "AABBB",
-    "i": "ABAAA",
-    "j": "BBBAA",
-    "k": "ABAAB",
-    "l": "ABABA",
-    "m": "ABABB",
-    "n": "ABBAA",
-    "o": "ABBAB",
-    "p": "ABBBA",
-    "q": "ABBBB",
-    "r": "BAAAA",
-    "s": "BAAAB",
-    "t": "BAABA",
-    "u": "BAABB",
-    "v": "BBBAB",
-    "w": "BABAA",
-    "x": "BABAB",
-    "y": "BABBA",
-    "z": "BABBB",
-    " ": " ",
+ENCODE_DICT = {
+    "a": "AAAAA", "b": "AAAAB", "c": "AAABA", "d": "AAABB",
+    "e": "AABAA", "f": "AABAB", "g": "AABBA", "h": "AABBB",
+    "i": "ABAAA", "j": "BBBAA", "k": "ABAAB", "l": "ABABA",
+    "m": "ABABB", "n": "ABBAA", "o": "ABBAB", "p": "ABBBA",
+    "q": "ABBBB", "r": "BAAAA", "s": "BAAAB", "t": "BAABA",
+    "u": "BAABB", "v": "BBBAB", "w": "BABAA", "x": "BABAB",
+    "y": "BABBA", "z": "BABBB", " ": " ",
 }
 
+DECODE_DICT = {value: key for key, value in ENCODE_DICT.items()}
 
-decode_dict = {value: key for key, value in encode_dict.items()}
 
-
-def encode(word: str) -> str:
+def encode(message: str, symbols: tuple[str, str] = ("A", "B")) -> str:
     """
-    Encodes to Baconian cipher
-
-    >>> encode("hello")
-    'AABBBAABAAABABAABABAABBAB'
-    >>> encode("hello world")
-    'AABBBAABAAABABAABABAABBAB BABAAABBABBAAAAABABAAAABB'
-    >>> encode("hello world!")
-    Traceback (most recent call last):
-        ...
-    Exception: encode() accepts only letters of the alphabet and spaces
+    Encode a message using Bacon's cipher.
+
+    Parameters
+    ----------
+    message : str
+        The message to encode (letters and spaces only).
+    symbols : tuple[str, str], optional
+        Custom symbols to replace A and B, by default ("A", "B").
+
+    Returns
+    -------
+    str
+        Encoded message using the chosen symbols.
+
+    Raises
+    ------
+    ValueError
+        If the message contains invalid characters.
+
+    Examples
+    --------
+    >>> encode("abc")
+    'AAAAA AAAAB AAABA'
+    >>> encode("abc", symbols=("X", "Y"))
+    'XXXXX XXXXY XXYXX'
+    >>> encode("hi there")
+    'AABBB ABAAA BAABA AABBB AABAA BAAAA AABAA'
     """
-    encoded = ""
-    for letter in word.lower():
-        if letter.isalpha() or letter == " ":
-            encoded += encode_dict[letter]
+    a_sym, b_sym = symbols
+    encoded_message = ""
+
+    for char in message.lower():
+        if char.isalpha() or char == " ":
+            bacon = ENCODE_DICT[char]
+            encoded_message += bacon.replace("A", a_sym).replace("B", b_sym)
         else:
-            raise Exception("encode() accepts only letters of the alphabet and spaces")
-    return encoded
+            raise ValueError("Message can only contain letters and spaces.")
+
+    return encoded_message
 
 
-def decode(coded: str) -> str:
+def decode(cipher: str, symbols: tuple[str, str] = ("A", "B")) -> str:
     """
-    Decodes from Baconian cipher
-
-    >>> decode("AABBBAABAAABABAABABAABBAB BABAAABBABBAAAAABABAAAABB")
-    'hello world'
-    >>> decode("AABBBAABAAABABAABABAABBAB")
-    'hello'
-    >>> decode("AABBBAABAAABABAABABAABBAB BABAAABBABBAAAAABABAAAABB!")
-    Traceback (most recent call last):
-        ...
-    Exception: decode() accepts only 'A', 'B' and spaces
+    Decode a Bacon's cipher message.
+
+    Parameters
+    ----------
+    cipher : str
+        The encoded message using two symbols.
+    symbols : tuple[str, str], optional
+        Symbols used in the cipher, by default ("A", "B").
+
+    Returns
+    -------
+    str
+        Decoded message.
+
+    Raises
+    ------
+    ValueError
+        If the cipher contains invalid symbols or cannot be decoded.
+
+    Examples
+    --------
+    >>> decode("AAAAA AAAAB AAABA")
+    'abc'
+    >>> decode("XXXXX XXXXY XXYXX", symbols=("X","Y"))
+    'abc'
     """
-    if set(coded) - {"A", "B", " "} != set():
-        raise Exception("decode() accepts only 'A', 'B' and spaces")
-    decoded = ""
-    for word in coded.split():
-        while len(word) != 0:
-            decoded += decode_dict[word[:5]]
-            word = word[5:]
-        decoded += " "
-    return decoded.strip()
+    sym1, sym2 = symbols
+    unique_symbols = set(cipher.replace(" ", ""))
+    if unique_symbols - {sym1, sym2}:
+        raise ValueError(f"Cipher must contain only symbols {sym1} and {sym2}.")
 
+    candidates = []
+    for mapping in [(sym1, sym2), (sym2, sym1)]:
+        s1, s2 = mapping
+        standard = cipher.replace(s1, "A").replace(s2, "B")
+        try:
+            decoded = ""
+            for word in standard.split():
+                while word:
+                    chunk = word[:5]
+                    if chunk not in DECODE_DICT:
+                        raise ValueError
+                    decoded += DECODE_DICT[chunk]
+                    word = word[5:]
+                decoded += " "
+            candidates.append(decoded.strip())
+        except ValueError:
+            candidates.append(None)
 
-if __name__ == "__main__":
-    from doctest import testmod
+    for candidate in candidates:
+        if candidate is not None:
+            return candidate
+
+    raise ValueError("No valid decoding found.")
+
+
+def detect_unique_symbols(cipher: str) -> tuple[str, str]:
+    """
+    Detects the two unique symbols used in a cipher.
+
+    Parameters
+    ----------
+    cipher : str
+        Encoded message containing exactly two unique symbols.
 
-    testmod()
+    Returns
+    -------
+    tuple[str, str]
+        The two unique symbols found in the cipher.
+
+    Raises
+    ------
+    ValueError
+        If cipher does not contain exactly two unique symbols.
+
+    Examples
+    --------
+    >>> detect_unique_symbols("XXXYX YXXYX")
+    ('X', 'Y')
+    """
+    letters_only = [char for char in set(cipher.replace(" ", "")) if char.isalpha()]
+    if len(letters_only) != 2:
+        raise ValueError("Cipher must contain exactly two unique alphabetic symbols.")
+    return tuple(letters_only)
+
+
+if __name__ == "__main__":
+    # Example usage
+    cipher_text = (
+        "FEEFE EEFFF EEFEE EFFFF FEEFF EFEEE EEEFE "
+        "EFEEF EEEEF FEEEE EFFEF FEFEE EFFEE EEFEF EFFEF FEFEF"
+    )
+    sym1, sym2 = detect_unique_symbols(cipher_text)
+    decoded_message = decode(cipher_text, symbols=(sym1, sym2))
+    print(decoded_message)  # Expected: 'the quick brown fox'