sqlglot.trie

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