Reddit mirror encryption challenge

While reviewing Encryption using a mirror field, I decided to write my own solution from scratch. The challenge, from Reddit's /r/dailyprogrammer, is to implement a monoalphabetic substitution cipher whose key is derived from a pictorial 13×13 grid. Here's a simplified 5×5 example:

 abcde
A \   f
B     g
C   \ h
D     i
E   / j
FGHIJ


The 5×5 grid above defines 20 mappings, some of which are:

• A maps to G, because a laser beam heading right from the A would hit the \ mirror and be deflected down to G.
• Conversely, G maps to A through the reverse path.
• B maps to g, and vice versa, as the horizontal beam is uninterrupted.
• C maps to E, and vice versa, as the beam is reflected twice.
• I maps to j, and vice versa.

The actual input is provided as a 13×13 grid of /, \, and space characters, without the row and column labels. I've implemented the option to accept the key as a filename if one is provided on the command line, else it will be read from the first 13 lines of standard input. Thereafter, the program acts as a filter for text on standard input.

I've also implemented a doctest, which you can run using python -v -m doctest programname. The presence of backslashes has forced me to write it in an ugly way; any suggestion for a better way to write it would be appreciated.

"""
Perform the mirror encryption challenge in https://redd.it/4m3ddb
"""

from fileinput import input as fileinput
from string import ascii_lowercase, ascii_uppercase
from sys import argv, stdin

def grid_to_translation(grid):
"""
Given a mirror field, construct a translation table suitable for use with
str.maketrans().  The grid may be any iterable; only the first 13 lines
will be taken from it.

>>> # Using a string replacement as literal backslashes are hard to write.
>>> input = (line.replace('x', chr(92)) for line in [
...     '   xx  /x    ',
...     '            x',
...     '   /         ',
...     '      x     x',
...     '    x        ',
...     '  /      /   ',
...     'x  /      x  ',
...     '     x       ',
...     'x/           ',
...     '/            ',
...     '          x  ',
...     '    x/       ',
...     '   /       / ',
... ])
>>> translation = grid_to_translation(input)
>>> 'TpnQSjdmZdpoohd'.translate(translation)
'DailyProgrammer'
>>> 'Code Review'.translate(translation)
'Amrh LhOnhN'
>>> 'Amrh LhOnhN'.translate(translation)
'Code Review'
"""
RIGHT, DOWN, LEFT, UP = (0, 1), (1, 0), (0, -1), (-1, 0)

LAUNCH = dict(      # Top, right, left, and bottom edges
[(c, ((-1, col), DOWN))  for col, c in enumerate(ascii_lowercase[:13])]
+ [(c, ((row, 13), LEFT))  for row, c in enumerate(ascii_lowercase[13:])]
+ [(c, ((row, -1), RIGHT)) for row, c in enumerate(ascii_uppercase[:13])]
+ [(c, ((13, col), UP))    for col, c in enumerate(ascii_uppercase[13:])]
)

EDGES = {pos: c for c, (pos, _) in LAUNCH.items()}

# REFLECT[barrier][incoming_direction] gives the outgoing direction
REFLECT = {
'/': {
DOWN: LEFT,
RIGHT: UP,      LEFT: DOWN,
UP: RIGHT,
},
'\\': {
DOWN: RIGHT,
RIGHT: DOWN,    LEFT: UP,
UP: LEFT,
},
' ': {
DOWN: DOWN,
RIGHT: RIGHT,   LEFT: LEFT,
UP: UP,
},
}

def trace(pos, aim):
pos = (pos[0] + aim[0], pos[1] + aim[1])
while not pos in EDGES:
aim = REFLECT[grid[pos[0]][pos[1]]][aim]
pos = (pos[0] + aim[0], pos[1] + aim[1])
return EDGES[pos]

grid = iter(grid)
grid = [next(grid) for _ in range(13)]
mapping = {ord(c): trace(pos, aim) for c, (pos, aim) in LAUNCH.items()}
assert all(mapping[ord(d)] == chr(c) for c, d in mapping.items())
return str.maketrans(mapping)

if __name__ == '__main__':
# Make a translation table from the a file specified on the command line,
# or from the first 13 lines of stdin if no filename is given.
translation = grid_to_translation(fileinput() if len(argv) > 1 else stdin)

# Then use it to encrypt all subsequent input lines.
for text in stdin:
print(text.translate(translation), end='')


• fileinput() already defaults to stdin if there is no arguments on the command line, no need to manage it yourself. It has been made clearer in a comment that you expect more than "no filename = stdin" : you should make it more clear in the code too; maybe changing the second line of comment with something like # since fileinput will truncate anything to the 13 lines we need, overide the 'no filename = stdin' rule to be able to pipe in the grid + some text to translate.
• You can easily use '\' in docstrings by using raw multiline strings: r'''docstring'''. Thus when defining input you can use '\\' as in a regular Python session, or use raw strings where applicable.
Or use raw strings all along and add padding characters for those that ends with a backslash, as you already ignore them; and document it with something along the lines of # avoid ending raw strings with backslashes, extra characters are ignored anyway.
• You can define the magic number 13 as MID_ALPHABET = len(ascii_lowercase) / 2
• Try it, though, and you may see the reasons for my workarounds. fileinput alone won't let me read the key from a file and the plaintext from stdin. Raw strings fail if the last character is a backslash. – 200_success Jun 11 '16 at 16:17
• When fileinput reads from stdin, either because there is no arguments or because there is a '-' on the command line, it doesn't process lines as you type it. You have to explicitly hit Ctrl+D to close the file before fileinput yields the lines. But you can still read from stdin later on. As long as you use fileinput.input() and thus accept that behaviour if the user enters - as first parameters, you might as well use it as the default. But I understand that your workaround lets you pipe the grid into the program and still access stdin later. – 409_Conflict Jun 11 '16 at 16:52