simple trie implementation

trie.py

# simple trie implementation

class Trie:
    class Node:
        def __init__(self, data=None, end=False):
            self.data = data
            self.children = dict()
            # is this the end of a word?
            self.end = end
        def __str__(self):
            return str(self.children)
        def __repr__(self):
            return str(self)
        def add_child(self, node):
            self.children[node.data] = node

    def __init__(self):
        self.root = Trie.Node()

    def __contains__(self, string):
        curr = self.root
        for i in range(len(string)):
            if string[i] in curr.children:
                curr = curr.children[string[i]]
            else:
                return False
        return True

    def __str__(self):
        return str(self.root)

    def __repr__(self):
        return str(self)

    def startswith(self, string):
        curr = self.root
        for i in range(len(string)):
            if string[i] in curr.children:
                curr = curr.children[string[i]]
            else:
                return set()

        ret = set()
        def traverse_children(node, current_string):
            if node.end:
                ret.add(current_string)
                return
            for child in node.children.keys():
                traverse_children(node.children[child], current_string + child)

        traverse_children(curr, string)
        return ret

    def add(self, string):
        curr = self.root
        for i in range(len(string)):
            if string[i] in curr.children:
                curr = curr.children[string[i]]
            else:
                break
        
        for j in range(i, len(string)):
            end = j == len(string) - 1
            curr.add_child(Trie.Node(string[j], end=end))
            curr = curr.children[string[j]]

if __name__ == '__main__':
    t = Trie()
    t.add('CAR')
    t.add('CAT')
    t.add('COB')

    print(t)
    
    print('COB' in t)
    print('COBD' in t)
    
    print(t.startswith('C'))