sqlglot.serde

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