Trie in Python

A post about Haskell vs. Python readability came onto my radar the other day. It compares two implementations of a trie structure, and after looking upon the Python version I wanted to make my own attempt. I didn't make it to necessarily compare or "battle" against the other solutions, it's more of an exercise in the vein of "how would I do it".

The code

Here's the original code (for easier lookup, as I refer to a few things in it in the notes):

class Trie(object):
    def __init__(self, value=None):
        self.children = {}
        self.value = value
        self.flag = False # Flag to represent that a word ends at this node

    def add(self, char):
        val = self.value + char if self.value else char
        self.children[char] = Trie(val)

    def insert(self, word):
        node = self
        for char in word:
            if char not in node.children:
                node.add(char)
            node = node.children[char]
        node.flag = True

    def find(self, word):
        node = self
        for char in word:
            if char not in node.children:
                return None
            node = node.children[char]
        return node.value

    def all_prefixes(self, wlist):
        results = set()
        if self.flag:
            results.add(self.value)
        if not self.children: return results
        return reduce(lambda a, b: a | b,
                     [node.all_prefixes() for
                      node in self.children.values()]) | results

    def autocomplete(self, prefix):
        node = self
        for char in prefix:
            if char not in node.children:
                return set()
            node = node.children[char]
        return node.all_prefixes()

My code:

class Trie:
    def __init__(self):
        self.children = {}
        self.is_word_end = False

    def insert(self, word):
        for char in word:
            self = self.children.setdefault(char, Trie())
        self.is_word_end = True

    def words_with(self, prefix):
        if self.is_word_end:
            yield prefix
        for char, node in self.children.items():
            yield from node.words_with(prefix + char)

    def autocomplete(self, prefix):
        try:
            for char in prefix:
                self = self.children[char]
            return list(self.words_with(prefix))
        except KeyError:
            return []

A few notes:

• Not storing self.value does seem to reduce complexity, perhaps counter-intuitively.
• The oft neglected dict.setdefault allowed me to inline the entire Trie.add.
• Another Pythonism, yield and yield from, is a nice pattern for recursive tree walking that would otherwise require temporary containers. It also usually results in tighter code.
• I attempted a couple more experiments with making the code more functional, like using reduce and recursions instead of for-loops, but it didn't improve things really. Python is not a functional language in its soul :-)
• The idea of re-binding self while tree-walking may scare some people, but I thought doing node = self just to avoid this was a bit silly :-)

P.S. Please don't talk to me about types and dataclasses, I will ridicule you to no end :-)