Edit on GitHub

# sqlglot.trie

View Source

     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)

key = typing.Sequence[typing.Hashable]
class TrieResult(enum.Enum): View Source

     8class TrieResult(Enum):
     9    FAILED = auto()
    10    PREFIX = auto()
    11    EXISTS = auto()

An enumeration.

FAILED = <TrieResult.FAILED: 1>
PREFIX = <TrieResult.PREFIX: 2>
EXISTS = <TrieResult.EXISTS: 3>

##### Inherited Members

enum.Enum
  name
  value

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

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

    44def in_trie(trie: t.Dict, key: key) -> t.Tuple[TrieResult, t.Dict]:
    45    """
    46    Checks whether a key is in a trie.
    47
    48    Examples:
    49        >>> in_trie(new_trie(["cat"]), "bob")
    50        (<TrieResult.FAILED: 1>, {'c': {'a': {'t': {0: True}}}})
    51
    52        >>> in_trie(new_trie(["cat"]), "ca")
    53        (<TrieResult.PREFIX: 2>, {'t': {0: True}})
    54
    55        >>> in_trie(new_trie(["cat"]), "cat")
    56        (<TrieResult.EXISTS: 3>, {0: True})
    57
    58    Args:
    59        trie: The trie to be searched.
    60        key: The target key.
    61
    62    Returns:
    63        A pair `(value, subtrie)`, where `subtrie` is the sub-trie we get at the point
    64        where the search stops, and `value` is a TrieResult value that can be one of:
    65
    66        - TrieResult.FAILED: the search was unsuccessful
    67        - TrieResult.PREFIX: `value` is a prefix of a keyword in `trie`
    68        - TrieResult.EXISTS: `key` exists in `trie`
    69    """
    70    if not key:
    71        return (TrieResult.FAILED, trie)
    72
    73    current = trie
    74    for char in key:
    75        if char not in current:
    76            return (TrieResult.FAILED, current)
    77        current = current[char]
    78
    79    if 0 in current:
    80        return (TrieResult.EXISTS, current)
    81
    82    return (TrieResult.PREFIX, current)

Checks whether a key is in a trie.

###### Examples:

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

>     >>> in_trie(new_trie(["cat"]), "ca")
>     (<TrieResult.PREFIX: 2>, {'t': {0: True}})

>     >>> in_trie(new_trie(["cat"]), "cat")
>     (<TrieResult.EXISTS: 3>, {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 a TrieResult value that can be one of:

> - TrieResult.FAILED: the search was unsuccessful
> - TrieResult.PREFIX: `value` is a prefix of a keyword in `trie`
> - TrieResult.EXISTS: `key` exists in `trie`