Edit on GitHub

# sqlglot.trie

View Source

     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)

def new_trie( keywords: Iterable[Sequence[Hashable]], trie: Optional[Dict] = None) -> Dict: View Source

     7def new_trie(keywords: t.Iterable[key], trie: t.Optional[t.Dict] = None) -> t.Dict:
     8    """
     9    Creates a new trie out of a collection of keywords.
    10
    11    The trie is represented as a sequence of nested dictionaries keyed by either single character
    12    strings, or by 0, which is used to designate that a keyword is in the trie.
    13
    14    Example:
    15        >>> new_trie(["bla", "foo", "blab"])
    16        {'b': {'l': {'a': {0: True, 'b': {0: True}}}}, 'f': {'o': {'o': {0: True}}}}
    17
    18    Args:
    19        keywords: the keywords to create the trie from.
    20        trie: a trie to mutate instead of creating a new one
    21
    22    Returns:
    23        The trie corresponding to `keywords`.
    24    """
    25    trie = {} if trie is None else trie
    26
    27    for key in keywords:
    28        current = trie
    29
    30        for char in key:
    31            current = current.setdefault(char, {})
    32        current[0] = True
    33
    34    return trie

Creates a new trie out of a collection of keywords.

The trie is represented as a sequence of nested dictionaries keyed by either single character strings, or by 0, which is used to designate that a keyword is in the trie.

###### Example:

>     >>> new_trie(["bla", "foo", "blab"])
>     {'b': {'l': {'a': {0: True, 'b': {0: True}}}}, 'f': {'o': {'o': {0: True}}}}

###### Arguments:

- keywords: the keywords to create the trie from.
- trie: a trie to mutate instead of creating a new one

###### Returns:

> The trie corresponding to `keywords`.

def in_trie(trie: Dict, key: Sequence[Hashable]) -> Tuple[int, Dict]: View Source

    37def in_trie(trie: t.Dict, key: key) -> t.Tuple[int, t.Dict]:
    38    """
    39    Checks whether a key is in a trie.
    40
    41    Examples:
    42        >>> in_trie(new_trie(["cat"]), "bob")
    43        (0, {'c': {'a': {'t': {0: True}}}})
    44
    45        >>> in_trie(new_trie(["cat"]), "ca")
    46        (1, {'t': {0: True}})
    47
    48        >>> in_trie(new_trie(["cat"]), "cat")
    49        (2, {0: True})
    50
    51    Args:
    52        trie: the trie to be searched.
    53        key: the target key.
    54
    55    Returns:
    56        A pair `(value, subtrie)`, where `subtrie` is the sub-trie we get at the point where the search stops, and `value`
    57        is either 0 (search was unsuccessful), 1 (`value` is a prefix of a keyword in `trie`) or 2 (`key is in `trie`).
    58    """
    59    if not key:
    60        return (0, trie)
    61
    62    current = trie
    63
    64    for char in key:
    65        if char not in current:
    66            return (0, current)
    67        current = current[char]
    68
    69    if 0 in current:
    70        return (2, current)
    71    return (1, current)

Checks whether a key is in a trie.

###### Examples:

>     >>> in_trie(new_trie(["cat"]), "bob")
>     (0, {'c': {'a': {'t': {0: True}}}})

>     >>> in_trie(new_trie(["cat"]), "ca")
>     (1, {'t': {0: True}})

>     >>> in_trie(new_trie(["cat"]), "cat")
>     (2, {0: True})

###### Arguments:

- trie: the trie to be searched.
- key: the target key.

###### Returns:

> A pair `(value, subtrie)`, where `subtrie` is the sub-trie we get at the point where the search stops, and `value` is either 0 (search was unsuccessful), 1 (`value` is a prefix of a keyword in `trie`) or 2 (`key is in`trie`).