very simple trie for learning Data Structures

trie.py

class Trie:
    trie = None

    END_WORD = '~END~'


    def __init__(self):
        self.trie = {}

    def add(self, word):
        current_dict = self.trie
        prefix = ''
        for letter in word:
            prefix += letter
            if prefix in current_dict:
                current_dict = current_dict[prefix]
            else:
                current_dict[prefix] = {}
                current_dict = current_dict[prefix]
        current_dict[self.END_WORD] = True

    def retrive_for_prefix(self, word):
        current_dict = self.trie
        prefix = ''
        for letter in word:
            prefix += letter
            if prefix in current_dict:
                current_dict = current_dict[prefix]
            else:
                return {}
        return {word: current_dict}

    def rebuild_words(self, trie, result = None):
        if result is None:
            result = []

        for el in trie:
            if trie[el] is True:
                continue

            if self.END_WORD in trie[el]:
                result.append(el)

            if isinstance(trie[el], dict):
                self.rebuild_words(trie[el], result)

        return result


trie = Trie()
trie.add("onforce")
trie.add("onthetable")
trie.add("oa")
trie.add("oaa")
trie.add("oaaaa")
trie.add("oabe")
trie.add("oace")
trie.add("oade")
trie.add("b")
trie.add("baida")
trie.add("b12")
trie.add("b15")

on_sub_tree = trie.retrive_for_prefix("on")
assert ["onforce", "onthetable"] == trie.rebuild_words(on_sub_tree)
oa_sub_tree = trie.retrive_for_prefix("oa")
assert 6 == len(trie.rebuild_words(oa_sub_tree))
ob1_sub_tree = trie.retrive_for_prefix("b1")
assert {"b12", "b15"} == {x for x in trie.rebuild_words(ob1_sub_tree)}