Python function to create a trie from a list of words.

build_trie.py

def create_trie(words, tokenize_fn=None, count_key='__count__', is_word_key='__is_word__'):
    top_node = {}

    if not tokenize_fn:
        tokenize_fn = lambda word: list(word)

    for word in words:
        cursor = top_node
        tokens = tokenize_fn(word)
        tokens_length = len(tokens)

        for index, token in enumerate(tokens):
            parent_node = cursor
            current_node = cursor.get(token, None)
            is_last_item = index == tokens_length - 1

            if current_node is None:
                current_node = {}
                cursor[token] = current_node

            if is_last_item:
                parent_node[count_key] = parent_node.get(count_key, 0) + 1
                current_node[is_word_key] = True

            cursor = current_node

    return top_node