Wildcard Dictionary Lookup

wildcard_lookup.py

class TRIE(object):
    """
    >>> "foo" in TRIE(["foo", "bar"])
    True
    >>> "fo." in TRIE(["foo", "bar"])
    True
    >>> "f.o" in TRIE(["foo", "bar"])
    True
    >>> ".oo" in TRIE(["foo", "bar"])
    True
    >>> "..o" in TRIE(["foo", "bar"])
    True
    >>> ".o." in TRIE(["foo", "bar"])
    True
    >>> "..o" in TRIE(["foo", "bar"])
    True
    >>> "bar" in TRIE(["foo", "bar"])
    True
    >>> ".ar" in TRIE(["foo", "bar"])
    True
    >>> "b.r" in TRIE(["foo", "bar"])
    True
    >>> "ba." in TRIE(["foo", "bar"])
    True
    >>> "b.." in TRIE(["foo", "bar"])
    True
    >>> ".a." in TRIE(["foo", "bar"])
    True
    >>> "..r" in TRIE(["foo", "bar"])
    True
    >>> "..." in TRIE(["foo", "bar"])
    True
    >>> "..t" in TRIE(["foo", "bar"])
    False
    >>> "baz" in TRIE(["foo", "bar"])
    False
    >>> ".." in TRIE(["foo", "bar"])
    False
    >>> "...." in TRIE(["foo", "bar"])
    False
    """

    end = '\0'

    def __init__(self, words):
        root = dict()
        for word in words:
            current_dict = root
            for c in word:
                current_dict = current_dict.setdefault(c, {})
            current_dict = current_dict.setdefault(TRIE.end, TRIE.end)
        self.trie = root

    def match(self, root, pattern, idx):
        if idx is len(pattern):
            return isinstance(root, dict) and TRIE.end in root
        c = pattern[idx]
        if c is '.':
            # wildcard
            for n in root:
                if self.match(root[n], pattern, idx + 1):
                    return True
            return False
        elif c in root:
            return self.match(root[c], pattern, idx + 1)
        else:
            return False

    def __contains__(self, pattern):
        return self.match(self.trie, pattern, 0)

if __name__ == '__main__':
    import doctest
    doctest.testmod()