Python 3 Playfair cipher encrypter and decrypter

Question

I have written a Python 3 program that can encrypt and decrypt Playfair cipher, though I know it has been implemented countless times, none of the scripts I have seen so far(e.g. https://codegolf.stackexchange.com/a/23363) is as performant as mine, and many of them only do the encryption part and not the decryption part, not to mention most of them are written in Python 2.

My code is self documentary, without any code obfuscation and nested for loops and nested dicts/lists and any other kinds of nonsense, it uses a simple flat list and divmod as basis of the transformations.

My code implements all the rules of Playfair cipher, along with some additional rules: there are a total of 600 bigrams defined by the rules of Playfair cipher, and 'XX' isn't one of them, however in practice 'XX' can be a very real possibility;

I have pulled a total of 105230 words from GCIDE, although only one of them contains 'xx' (vioxx), 144 of them starts with x and 400 of them ends with x, there are three cases where xx may occur:

• case 1: a word not found in dictionaries contains xx

• case 2: a word that ends with x followed by a word that starts with x

• case 3: the last letter in the message is x and the message contains odd number of letters

When one of the above is true, standard Playfair cipher can't encrypt xx, but in my algorithm this is just a special case of rule 2.

Additionally, because of rule 1, there is a need to remove every x in the second position in a bigram, however bigram with x as its second letter can very possibly occur in a word, and when this happens an additional x should be inserted between the letters to break the bigram, for example, 'axe' after preparation is 'AX XE'.

The code:

import re
from enum import Enum


LETTERS = 'ABCDEFGHIKLMNOPQRSTUVWXYZ'


class Playfair:
    @staticmethod
    def check_input(arg):
        if not isinstance(arg, str):
            raise TypeError(
                'The argument of the function should be an instance of `str`')

        if not all(i.isascii() & i.isprintable() for i in arg):
            raise ValueError(
                'The argument of the function should be a `str` containing only printable ASCII characters')

    @staticmethod
    def prepare(msg):
        msg = msg.upper().replace('J', 'I')
        Playfair.check_input(msg)
        msg = ''.join([i for i in msg if i.isalpha()])
        while msg:
            s = msg[:2]
            if len(set(s)) == 2 and s[-1] != 'X':
                msg = msg[2:]
            else:
                s = s[0] + 'X'
                msg = msg[1:]
            yield s

    def __init__(self, key):
        key = key.upper().replace('J', 'I')
        Playfair.check_input(key)
        key = ''.join(i for i in key if i.isalpha())
        self.square = list(dict.fromkeys(key + LETTERS))

    def locate(self, letter):
        letter = letter.upper()
        return list(divmod(self.square.index(letter), 5))

    def get_letter(self, pos):
        index = pos[0] * 5 + pos[1]
        return self.square[index]

    def transform(self, pair, move):
        p0, p1 = map(self.locate, pair)
        def roll(x): return (x + move) % 5
        for i in range(3):
            if i < 2 and p0[i] == p1[i]:
                j = 1 - i
                p0[j], p1[j] = map(roll, (p0[j], p1[j]))
                break
            if i == 2:
                p0[1], p1[1] = p1[1], p0[1]
        return ''.join(map(self.get_letter, (p0, p1)))

    def encrypt(self, msg):
        pairs = Playfair.prepare(msg)
        return ' '.join(self.transform(pair, 1) for pair in pairs)

    def decrypt(self, msg):
        msg = msg.upper()
        Playfair.check_input(msg)
        if not re.match(r'^([A-Z]{2}(\s)?)+$', msg):
            raise ValueError(
                "The inputted message isn't a valid Playfair cipher message")
        msg = re.findall('[A-Z]{2}', msg)
        def rule1(x): return x if x[1] != 'X' else x[0]
        return ' '.join(rule1(self.transform(s, -1)) for s in msg)


def test():
    p = Playfair("Imagine there's no heaven")
    plain = [
        'I am a new soul, I came to this strange world',
        'Starry starry night, paint your palette blue and grey',
        'I crown me king of the sweet cold north',
        'If I die young, bury me in satin',
        "Come gather 'round people, wherever you roam",
        'Somewhere over the rainbow'
    ]
    encrypted = [
        'MG AG IS ZT FI FA BG IT EV HR NE EH HG IN TU CE QB',
        'EH GH CG EH GH CG IM AR RK GM MS UC YE LG FH HW HT LX IO GI CN ST ZY',
        'GO EC ZM IT FM IN FU HR TE UT TH DV QB ID SH BA',
        'EU NO EO UC ZI AC YE WG OE SD MH MI',
        'DV IT NG HR TS EC ZI CQ OF QP TU RT ST OT CG FI EC GA',
        'ED IT XT TS OF OT SH RT HG MI CV XY'
    ]
    decrypted = [
        'IA MA NE WS OU LI CA ME TO TH IS ST RA NG EW OR LD',
        'ST AR RY ST AR RY NI GH TP AI NT YO UR PA LE T TE BL UE AN DG RE Y',
        'IC RO WN ME KI NG OF TH ES WE ET CO LD NO RT H',
        'IF ID IE YO UN GB UR YM EI NS AT IN',
        'CO ME GA TH ER RO UN DP EO PL EW HE RE VE RY OU RO AM',
        'SO ME WH ER EO VE RT HE RA IN BO W'
    ]
    assert [p.encrypt(i) for i in plain] == encrypted
    assert [p.decrypt(i) for i in encrypted] == decrypted


if __name__ == '__main__':
    test()

How can it be improved?

P.S. I have tried to pregenerate the bigrams using itertools.permutations (naturally xx isn't supported by this method) and corresponding encrypted bigrams, to generate two lookup tables for encryption and decryption.

I have written two versions of this, one uses two dicts, the other uses two pandas.Series, and with the rest of the code unchanged, and ran the test function.

Surprisingly, the one without pregeneration is the fastest, %timeit states it takes around 1 ms on average to complete, and the one using the dicts is in the middle, uses 2.5 ms on average to complete, and the one using vectorized code is the slowest, it takes around 8.2 ms on average to complete...

Answer 1

First of all

• This code looks like it does what you want
• You've separated methods well
• The code style looks reasonable (in terms of formatting, blank lines etc).

Second, stop reading code golf, it's a poor standard even if you're trying to learn what not to do.

Third, some specific improvements I might suggest:

1. No, this code is NOT fully self-documenting. Add documentation. Add an explanation at the top of what Playfair is (a wikipedia link would be OK). Improve your method and variable names. Add a docstring for each method explaining what it does (returns) -- it's not clear enough what transform, get_letter, locate, or check_input do based on the name alone. encrypt and decrypt are fine, although (see below) encrypt does more than one might expect without a docstring.
2. Minor, but it's considered poor style to define one-line inline functions. Move roll and _rule to the top level of the class. If these are only used by that method, a common way to indicate this might be to put no blank lines between them, and call the methods _roll and _rule1.
3. Pull out the logic which transforms the input plaintext into bigraphs from the encrypt function, into its own helper function. Update the test function accordingly.
4. Pull out the logic to avoid "XX" into its own function. Right now this is a lot of your explanation and your code--consider removing it or simplifying it, to make the code cleaner at the expense of covering all cases. This doesn't seem worth the complexity.
5. Consider having decrypt remove spaces in the output. You might also consider returning all-lowercase as being more readable to modern viewers than all-caps. The main benefit to capital letter is avoiding legibility errors when performing ciphers by hand.
6. In test, I would recommend grouping together the plaintext, encrypted, and decrypted forms of each phrase, rather than having three parallel lists.
7. Most python readers prefer list comprehensions to map. It's up to you if you care.
8. Remove the for loop in transform. This logic is confusing to read. Use three flat cases instead.