sqlglot.serde API documentation

Edit on GitHub

# sqlglot.serde

View Source

      1from __future__ import annotations
      2
      3import typing as t
      4
      5from sqlglot import expressions as exp
      6
      7
      8INDEX = "i"
      9ARG_KEY = "k"
     10IS_ARRAY = "a"
     11CLASS = "c"
     12TYPE = "t"
     13COMMENTS = "o"
     14META = "m"
     15VALUE = "v"
     16DATA_TYPE = "DataType.Type"
     17
     18
     19def dump(expression: exp.Expression) -> t.List[t.Dict[str, t.Any]]:
     20    """
     21    Dump an Expression into a JSON serializable List.
     22    """
     23    i = 0
     24    payloads = []
     25    stack: t.List[t.Tuple[t.Any, t.Optional[int], t.Optional[str], bool]] = [
     26        (expression, None, None, False)
     27    ]
     28
     29    while stack:
     30        node, index, arg_key, is_array = stack.pop()
     31
     32        payload: t.Dict[str, t.Any] = {}
     33
     34        if index is not None:
     35            payload[INDEX] = index
     36        if arg_key is not None:
     37            payload[ARG_KEY] = arg_key
     38        if is_array:
     39            payload[IS_ARRAY] = is_array
     40
     41        payloads.append(payload)
     42
     43        if hasattr(node, "parent"):
     44            klass = node.__class__.__qualname__
     45
     46            if node.__class__.__module__ != exp.__name__:
     47                klass = f"{node.__module__}.{klass}"
     48
     49            payload[CLASS] = klass
     50
     51            if node.type:
     52                payload[TYPE] = dump(node.type)
     53            if node.comments:
     54                payload[COMMENTS] = node.comments
     55            if node._meta is not None:
     56                payload[META] = node._meta
     57            if node.args:
     58                for k, vs in reversed(node.args.items()):
     59                    if type(vs) is list:
     60                        for v in reversed(vs):
     61                            stack.append((v, i, k, True))
     62                    elif vs is not None:
     63                        stack.append((vs, i, k, False))
     64        elif type(node) is exp.DataType.Type:
     65            payload[CLASS] = DATA_TYPE
     66            payload[VALUE] = node.value
     67        else:
     68            payload[VALUE] = node
     69
     70        i += 1
     71
     72    return payloads
     73
     74
     75@t.overload
     76def load(payloads: None) -> None: ...
     77
     78
     79@t.overload
     80def load(payloads: t.List[t.Dict[str, t.Any]]) -> exp.Expression: ...
     81
     82
     83def load(payloads):
     84    """
     85    Load a list of dicts generated by dump into an Expression.
     86    """
     87
     88    if not payloads:
     89        return None
     90
     91    payload, *tail = payloads
     92    root = _load(payload)
     93    nodes = [root]
     94    for payload in tail:
     95        node = _load(payload)
     96        nodes.append(node)
     97        parent = nodes[payload[INDEX]]
     98        arg_key = payload[ARG_KEY]
     99
    100        if payload.get(IS_ARRAY):
    101            parent.append(arg_key, node)
    102        else:
    103            parent.set(arg_key, node)
    104
    105    return root
    106
    107
    108def _load(payload: t.Dict[str, t.Any]) -> exp.Expression | exp.DataType.Type:
    109    class_name = payload.get(CLASS)
    110
    111    if not class_name:
    112        return payload[VALUE]
    113    if class_name == DATA_TYPE:
    114        return exp.DataType.Type(payload[VALUE])
    115
    116    if "." in class_name:
    117        module_path, class_name = class_name.rsplit(".", maxsplit=1)
    118        module = __import__(module_path, fromlist=[class_name])
    119    else:
    120        module = exp
    121
    122    expression = getattr(module, class_name)()
    123    expression.type = load(payload.get(TYPE))
    124    expression.comments = payload.get(COMMENTS)
    125    expression._meta = payload.get(META)
    126    return expression

INDEX = 'i'
ARG_KEY = 'k'
IS_ARRAY = 'a'
CLASS = 'c'
TYPE = 't'
COMMENTS = 'o'
META = 'm'
VALUE = 'v'
DATA_TYPE = 'DataType.Type'
def dump(expression: sqlglot.expressions.Expression) -> List[Dict[str, Any]]: View Source

    20def dump(expression: exp.Expression) -> t.List[t.Dict[str, t.Any]]:
    21    """
    22    Dump an Expression into a JSON serializable List.
    23    """
    24    i = 0
    25    payloads = []
    26    stack: t.List[t.Tuple[t.Any, t.Optional[int], t.Optional[str], bool]] = [
    27        (expression, None, None, False)
    28    ]
    29
    30    while stack:
    31        node, index, arg_key, is_array = stack.pop()
    32
    33        payload: t.Dict[str, t.Any] = {}
    34
    35        if index is not None:
    36            payload[INDEX] = index
    37        if arg_key is not None:
    38            payload[ARG_KEY] = arg_key
    39        if is_array:
    40            payload[IS_ARRAY] = is_array
    41
    42        payloads.append(payload)
    43
    44        if hasattr(node, "parent"):
    45            klass = node.__class__.__qualname__
    46
    47            if node.__class__.__module__ != exp.__name__:
    48                klass = f"{node.__module__}.{klass}"
    49
    50            payload[CLASS] = klass
    51
    52            if node.type:
    53                payload[TYPE] = dump(node.type)
    54            if node.comments:
    55                payload[COMMENTS] = node.comments
    56            if node._meta is not None:
    57                payload[META] = node._meta
    58            if node.args:
    59                for k, vs in reversed(node.args.items()):
    60                    if type(vs) is list:
    61                        for v in reversed(vs):
    62                            stack.append((v, i, k, True))
    63                    elif vs is not None:
    64                        stack.append((vs, i, k, False))
    65        elif type(node) is exp.DataType.Type:
    66            payload[CLASS] = DATA_TYPE
    67            payload[VALUE] = node.value
    68        else:
    69            payload[VALUE] = node
    70
    71        i += 1
    72
    73    return payloads

Dump an Expression into a JSON serializable List.

def load(payloads): View Source

     84def load(payloads):
     85    """
     86    Load a list of dicts generated by dump into an Expression.
     87    """
     88
     89    if not payloads:
     90        return None
     91
     92    payload, *tail = payloads
     93    root = _load(payload)
     94    nodes = [root]
     95    for payload in tail:
     96        node = _load(payload)
     97        nodes.append(node)
     98        parent = nodes[payload[INDEX]]
     99        arg_key = payload[ARG_KEY]
    100
    101        if payload.get(IS_ARRAY):
    102            parent.append(arg_key, node)
    103        else:
    104            parent.set(arg_key, node)
    105
    106    return root

Load a list of dicts generated by dump into an Expression.