sqlglot.serde

  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
 75def load(
 76    payloads: t.Optional[t.List[t.Dict[str, t.Any]]],
 77) -> t.Optional[exp.Expression | exp.DataType.Type]:
 78    """
 79    Load a list of dicts generated by dump into an Expression.
 80    """
 81
 82    if not payloads:
 83        return None
 84
 85    payload, *tail = payloads
 86    root = _load(payload)
 87    nodes = [root]
 88    for payload in tail:
 89        node = _load(payload)
 90        nodes.append(node)
 91        parent = nodes[payload[INDEX]]
 92        arg_key = payload[ARG_KEY]
 93
 94        if payload.get(IS_ARRAY):
 95            parent.append(arg_key, node)
 96        else:
 97            parent.set(arg_key, node)
 98
 99    return root
100
101
102def _load(payload: t.Dict[str, t.Any]) -> exp.Expression | exp.DataType.Type:
103    class_name = payload.get(CLASS)
104
105    if not class_name:
106        return payload[VALUE]
107    if class_name == DATA_TYPE:
108        return exp.DataType.Type(payload[VALUE])
109
110    if "." in class_name:
111        module_path, class_name = class_name.rsplit(".", maxsplit=1)
112        module = __import__(module_path, fromlist=[class_name])
113    else:
114        module = exp
115
116    expression = getattr(module, class_name)()
117    expression._type = load(payload.get(TYPE))
118    expression.comments = payload.get(COMMENTS)
119    expression._meta = payload.get(META)
120    return expression