A recursive word square generator

word_squares.py

#!/usr/bin/env python3

''' Generate word squares

    Written by PM 2Ring 1999.01.18
    Translated from C to Python 2017.01.20
'''

import sys
from bisect import bisect_left

# Default word list file name. 
# The list must be in ascending order.
# Words must be one per line, upper case, 
# with no whitespace (apart from the newline).
# Sowpods wordlist from http://www.3zsoftware.com/download/
WORD_FILE = 'scrabble_wordlist_sowpods.txt'

def insert_word(grid, i, s):
    ''' Copy word `s` into row & column `i` of `grid`  '''
    for j in range(i, WORD_LEN):
        grid[i][j] = grid[j][i] = s[j]

def show(grid):
    print('\n'.join([' '.join(row) for row in grid]), end='\n\n')

def do_square(grid, n):
    ''' Recursively complete the word square in `grid` '''

    # Get the first `n` chars of the word at row / column `n` of `grid`
    word = ''.join(grid[n][:n])

    # Make the next string (in alphabetical order) after `word`
    nextword = word[:-1] + chr(ord(word[-1]) + 1)

    # Find the range of indices in WORDS of all the words
    # that match the first `n` chars of `word`
    start = bisect_left(WORDS, word)
    stop = bisect_left(WORDS, nextword)

    # Insert the matching words into the grid
    for s in WORDS[start:stop]:
        insert_word(grid, n, s)
        if n == WORD_LEN - 1:
            yield grid
        else:
            yield from do_square(grid, n + 1)


if __name__ == '__main__':
    if '-h' in sys.argv or len(sys.argv) < 2:
        fmt = 'Generate word squares.\nUsage: %s seed_word [word_file]'
        print(fmt % sys.argv[0])
        sys.exit()
    seed_word = sys.argv[1].upper()
    word_file = sys.argv[2] if len(sys.argv) > 2 else WORD_FILE

    WORD_LEN = len(seed_word)

    # Build a list of the words we may need
    letters = set(seed_word)
    tlen = WORD_LEN + 1
    with open(word_file) as f:
        WORDS = [word[:-1] for word in f 
            if len(word) == tlen and word[0] in letters]

    # Create an empty grid & insert the seed word
    grid = [[''] * WORD_LEN for _ in range(WORD_LEN)]
    insert_word(grid, 0, seed_word)

    # Find & print all the solutions
    for i, g in enumerate(do_square(grid, 1), 1):
        print(i)
        show(g)