Autocomplete with a trie

autocomplete.py

#!/usr/bin/python3
'''
Autocompletion using a trie and a user supplied language file.
HINT: Linux/Unix users can look in "/usr/share/dict" for language files.

Author: rwalker
Email: r_walker@zoho.com
License: MIT
'''
import argparse, os, sys

class Trie:
    '''A simple tree structure to handle looking up all words starting
       with a given prefix'''

    def __init__(self, achar):
        self.char = achar
        self.children = {}
        
    def insert_word(self, aword):
        '''load a word into the trie'''
        cur = self
        for achar in aword:
            nxt = cur.children.get(achar)
            if nxt is None:
                nxt = Trie(achar)
                cur.children[achar] = nxt
            cur = nxt

def walk(trie,s):
    '''Walk the trie to the end of the prefix s'''
    cur = trie
    for c in s:
        cur = cur.children.get(c)
        if cur is None:
            return cur
    return cur

def print_all_words(trie, s = ""):
    '''print all the words in this trie with prefix s appended.'''
    if trie is None:
        return
    if len(trie.children)==0:
        print(s)
    else:
        for k,v in trie.children.items():
            print_all_words(v,s+v.char)

def autocomplete(root, s):
    '''An autocomplete method to lookup s in a trie'''
    start = walk(root, s)
    print_all_words(start, s)
    
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Suggest word completions based on a dictionary.")
    parser.add_argument("D", help="A dictionary file (one word per line)")
    args = parser.parse_args()
    
    # read the dictionary into the trie
    with open(args.D, encoding='utf-8', errors='ignore') as f:
        words = f.readlines()
    
    root = Trie("")
    for word in words:
        root.insert_word(word.strip("\n")+"\0")

    # take in strings from the command line
    try:
        while True:
            query = input("Prefix string: ")
            autocomplete(root, query)
    except KeyboardInterrupt:
        try:
            sys.exit(0)
        except SystemExit:
            os._exit(0)