Prefix tree implementation in Python 3, based on `dict`

trie.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
from typing import Sequence, Union
# Does not typecheck yet because of the hardcoded slice int types

class KeyErrorWithChildren(KeyError):
    pass

TrieKey = Union[Sequence,slice]

class Trie(dict):
    """
    Trie implementation:

    >>> t = Trie()
    >>> t['a'] = 3
    >>> t
    Trie('a': Trie(3))
    >>> t['abc'] = 5
    >>> t
    Trie('a': Trie(3, 'b': Trie('c': Trie(5))))

    When accessing a key without a value, the outcomes are:

    If the node has childrens, the exception KeyErrorWithChildren is raised:

    >>> t['ab']
    Traceback (most recent call last):
    ...
    trie.KeyErrorWithChildren: 'ab'

    It is possible to access the children accessing the `slice` syntax:

    >>> t['ab':]
    Trie('c': Trie(5))

    Instead if a node is empty and has no children, a regular KeyError is raised.

    >>> t['abd']
    Traceback (most recent call last):
    ...
    KeyError: 'abd'

    The Trie can work at a bigger "resolution" too:

    >>> t = Trie()
    >>> t['/usr/bin/git'.split('/')] = 12345
    >>> t
    Trie('': Trie('usr': Trie('bin': Trie('git': Trie(12345)))))
    """
    __slots__ = ('value', )
    class Empty:
        pass
    def __init__(self):
        self.value = self.Empty
    def _get_branch(self, key: Sequence):
        branch = self
        if key is not None:
            for part in key:
                if part in branch:
                    branch = super(Trie, branch).__getitem__(part)
                else:
                    child = Trie()
                    super(Trie, branch).__setitem__(part, child)
                    branch = child
        return branch
    def __getitem__(self, key: TrieKey):
        if isinstance(key, slice):
            branch = self._get_branch(key.start)
            children = True
        else:
            branch = self._get_branch(key)
            children = False
        if children:
            return branch
        else:
            if branch.value is self.Empty:
                if len(branch):
                    raise KeyErrorWithChildren(key)
                raise KeyError(key)
            return branch.value
    def __setitem__(self, key: TrieKey, value):
        branch = self._get_branch(key)
        branch.value = value
    def __repr__(self):
        args= []  # type: List[str]
        if self.value is not self.Empty:
            args += [repr(self.value)]
        if len(self):
            args += [", ".join("%r: %r" % (k, v) for k, v in self.items())]
        return "Trie(%s)" % ", ".join(args)