Affine cypher (exercism solution)

affine_cipher.py

import string
from itertools import count, cycle
from math import gcd

alphabet = string.ascii_lowercase
M = len(alphabet) # Length of roman alphabet
GROUP_LENGTH = 5
sep = [''] * (GROUP_LENGTH - 1) + [' ']

def mmi(a, m):
    """
    Calculate the modular multiplicative inverse of a related to m. 

    Parameters
    ----------
    a : int
        DESCRIPTION.
    m : int
        DESCRIPTION.

    Raises
    ------
    ValueError
        Raises if a and m aren't coprimes.

    Returns
    -------
    int
        Modular multiplicative inverse of a related to m.

    """
    if gcd(a, m) != 1:
        raise ValueError("a and m must be coprimes")
    m26plus1 = ((i * m + 1) for i in count())
    a_mod_m = a % m
    for attempt in m26plus1:
        if attempt % a_mod_m == 0:
            break
    return attempt // a_mod_m

def encode(plain_text, a, b):
    """
    Encode plaint_text using the affine cypher with the keys a and b.

    Parameters
    ----------
    plain_text : str
        Text to be encoded.
    a : int
        Multiplicative key.
    b : int
        Additive key.

    Raises
    ------
    ValueError
        Raises if a and the alphabet length aren't coprimes..

    Returns
    -------
    str
        Encoded str in batches of 5.

    """
    if gcd(a, M) != 1:
        raise ValueError("a and m must be coprimes")
    group_sep = (ch for ch in cycle(sep))
    codes = [
        ch if ch.isdigit() else (a * alphabet.find(ch.lower()) + b) % M
        for ch in plain_text
        if ch.lower() in alphabet or ch.isdigit()
    ]
    # return ''.join([f"{alphabet[code]}{' ' if (i + 1) % 5 == 0 else ''}" for i, code in enumerate(codes)])
    return ''.join([
        f"{code if type(code) == str else alphabet[code]}{next(group_sep)}"
        for code in codes
    ]).rstrip()


def decode(ciphered_text, a, b):
    """
    Decode ciphered_text using the affine cypher with the keys a and b.

    Parameters
    ----------
    ciphered_text : str
        Encoded text.
    a : int
        Multiplicative key.
    b : int
        Addtive key.

    Returns
    -------
    str
        Decoded ciphered_text.

    """
    a_inverse = mmi(a, M)
    codes = [
        ch if ch.isdigit() else a_inverse * (alphabet.find(ch.lower()) - b) % M
        for ch in ciphered_text
        if ch.lower() in alphabet or ch.isdigit()
    ]
    return ''.join([f"{code if type(code) == str else alphabet[code]}" for code in codes])

affine_cipher_test.py

import unittest

from affine_cipher import decode, encode

# Tests adapted from `problem-specifications//canonical-data.json`


class AffineCipherTest(unittest.TestCase):
    def test_encode_yes(self):
        self.assertEqual(encode("yes", 5, 7), "xbt")

    def test_encode_no(self):
        self.assertEqual(encode("no", 15, 18), "fu")

    def test_encode_omg(self):
        self.assertEqual(encode("OMG", 21, 3), "lvz")

    def test_encode_o_m_g(self):
        self.assertEqual(encode("O M G", 25, 47), "hjp")

    def test_encode_mindblowingly(self):
        self.assertEqual(encode("mindblowingly", 11, 15), "rzcwa gnxzc dgt")

    def test_encode_numbers(self):
        self.assertEqual(
            encode("Testing,1 2 3, testing.", 3, 4), "jqgjc rw123 jqgjc rw"
        )

    def test_encode_deep_thought(self):
        self.assertEqual(encode("Truth is fiction.", 5, 17), "iynia fdqfb ifje")

    def test_encode_all_the_letters(self):
        self.assertEqual(
            encode("The quick brown fox jumps over the lazy dog.", 17, 33),
            "swxtj npvyk lruol iejdc blaxk swxmh qzglf",
        )

    def test_encode_with_a_not_coprime_to_m(self):
        with self.assertRaisesWithMessage(ValueError):
            encode("This is a test.", 6, 17)

    def test_decode_exercism(self):
        self.assertEqual(decode("tytgn fjr", 3, 7), "exercism")

    def test_decode_a_sentence(self):
        self.assertEqual(
            decode("qdwju nqcro muwhn odqun oppmd aunwd o", 19, 16),
            "anobstacleisoftenasteppingstone",
        )

    def test_decode_numbers(self):
        self.assertEqual(decode("odpoz ub123 odpoz ub", 25, 7), "testing123testing")

    def test_decode_all_the_letters(self):
        self.assertEqual(
            decode("swxtj npvyk lruol iejdc blaxk swxmh qzglf", 17, 33),
            "thequickbrownfoxjumpsoverthelazydog",
        )

    def test_decode_with_no_spaces_in_input(self):
        self.assertEqual(
            decode("swxtjnpvyklruoliejdcblaxkswxmhqzglf", 17, 33),
            "thequickbrownfoxjumpsoverthelazydog",
        )

    def test_decode_with_too_many_spaces(self):
        self.assertEqual(
            decode("vszzm    cly   yd cg    qdp", 15, 16), "jollygreengiant"
        )

    def test_decode_with_a_not_coprime_to_m(self):
        with self.assertRaisesWithMessage(ValueError):
            decode("Test", 13, 5)

    # Utility functions
    def assertRaisesWithMessage(self, exception):
        return self.assertRaisesRegex(exception, r".+")


if __name__ == "__main__":
    unittest.main()