sqlglot.trie

 1import typing as t
 2from enum import Enum, auto
 3
 4key = t.Sequence[t.Hashable]
 5
 6
 7class TrieResult(Enum):
 8    FAILED = auto()
 9    PREFIX = auto()
10    EXISTS = auto()
11
12
13def new_trie(keywords: t.Iterable[key], trie: t.Optional[t.Dict] = None) -> t.Dict:
14    """
15    Creates a new trie out of a collection of keywords.
16
17    The trie is represented as a sequence of nested dictionaries keyed by either single
18    character strings, or by 0, which is used to designate that a keyword is in the trie.
19
20    Example:
21        >>> new_trie(["bla", "foo", "blab"])
22        {'b': {'l': {'a': {0: True, 'b': {0: True}}}}, 'f': {'o': {'o': {0: True}}}}
23
24    Args:
25        keywords: the keywords to create the trie from.
26        trie: a trie to mutate instead of creating a new one
27
28    Returns:
29        The trie corresponding to `keywords`.
30    """
31    trie = {} if trie is None else trie
32
33    for key in keywords:
34        current = trie
35        for char in key:
36            current = current.setdefault(char, {})
37
38        current[0] = True
39
40    return trie
41
42
43def in_trie(trie: t.Dict, key: key) -> t.Tuple[TrieResult, t.Dict]:
44    """
45    Checks whether a key is in a trie.
46
47    Examples:
48        >>> in_trie(new_trie(["cat"]), "bob")
49        (<TrieResult.FAILED: 1>, {'c': {'a': {'t': {0: True}}}})
50
51        >>> in_trie(new_trie(["cat"]), "ca")
52        (<TrieResult.PREFIX: 2>, {'t': {0: True}})
53
54        >>> in_trie(new_trie(["cat"]), "cat")
55        (<TrieResult.EXISTS: 3>, {0: True})
56
57    Args:
58        trie: The trie to be searched.
59        key: The target key.
60
61    Returns:
62        A pair `(value, subtrie)`, where `subtrie` is the sub-trie we get at the point
63        where the search stops, and `value` is a TrieResult value that can be one of:
64
65        - TrieResult.FAILED: the search was unsuccessful
66        - TrieResult.PREFIX: `value` is a prefix of a keyword in `trie`
67        - TrieResult.EXISTS: `key` exists in `trie`
68    """
69    if not key:
70        return (TrieResult.FAILED, trie)
71
72    current = trie
73    for char in key:
74        if char not in current:
75            return (TrieResult.FAILED, current)
76        current = current[char]
77
78    if 0 in current:
79        return (TrieResult.EXISTS, current)
80
81    return (TrieResult.PREFIX, current)