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
 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