Find all words beginning with a given prefix

Question

Today I present some code for finding all words that begin with a given prefix and some simple test cases for review.

from itertools import chain


sentinel = object()


def insert(tree, word):
    for letter in chain(word, [sentinel]):
        tree = tree.setdefault(letter, {})


def words_beginning_with(start, tree):
    for letter in start:
        try:
            tree = tree[letter]
        except KeyError:
            return
    for word in _all_words(tree):
        yield start + word


def _all_words(tree, start=''):
    for letter, tree in tree.items():
        if letter is sentinel:
            yield start
        else:
            for word in _all_words(tree, start+letter):
                yield word


def test_tree():
    # Simple tests
    tree = {}

    assert [] == list(words_beginning_with('foo', tree))

    insert(tree, 'foo')
    assert {'f': {'o': {'o': {sentinel: {}}}}} == tree

    assert ['foo'] == list(words_beginning_with('foo', tree))

    insert(tree, 'food')
    assert {'f': {'o': {'o': {sentinel: {},
                              'd': {sentinel: {}}}}}} == tree

    assert sorted(['foo', 'food']) == sorted(words_beginning_with('fo', tree))

    insert(tree, 'foody')
    assert sorted(['foo', 'food', 'foody']) == sorted(
        words_beginning_with('fo', tree))

    insert(tree, 'fold')
    assert sorted(['foo', 'fold', 'food', 'foody']) == sorted(
        words_beginning_with('fo', tree))


if __name__ == "__main__":
    test_tree()

I appreciate brutal and honest feedback the most, thanks.

Answer 1

• Sentinel could simply be None instead of a "custom" object.
• You don't need chain either: just set the sentinel manually at the end of the for loop.
• If using Python 3, you can use yield from ... instead of from ...: yield.

def insert(tree, word):
    for letter in word:
        tree = tree.setdefault(letter, {})
    tree[None] = {}  # or = None, it doesn't matter


def words_beginning_with(start, tree):
    for letter in start:
        try:
            tree = tree[letter]
        except KeyError:
            return
    for word in _all_words(tree):
        yield start + word


def _all_words(tree, start=''):
    for letter, tree in tree.items():
        if letter is None:
            yield start
        else:
            yield from _all_words(tree, start + letter):