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.Expr) -> list[dict[str, t.Any]]:
 20    """
 21    Dump an Expr into a JSON serializable List.
 22    """
 23    i = 0
 24    payloads = []
 25    stack: list[tuple[t.Any, int | None, str | None, bool]] = [(expression, None, None, False)]
 26
 27    while stack:
 28        node, index, arg_key, is_array = stack.pop()
 29
 30        payload: dict[str, t.Any] = {}
 31
 32        if index is not None:
 33            payload[INDEX] = index
 34        if arg_key is not None:
 35            payload[ARG_KEY] = arg_key
 36        if is_array:
 37            payload[IS_ARRAY] = is_array
 38
 39        payloads.append(payload)
 40
 41        if hasattr(node, "parent"):
 42            klass = node.__class__.__qualname__
 43
 44            if node.__class__.__module__ != exp.__name__:
 45                klass = f"{node.__module__}.{klass}"
 46
 47            payload[CLASS] = klass
 48
 49            if node.type:
 50                payload[TYPE] = dump(node.type)
 51            if node.comments:
 52                payload[COMMENTS] = node.comments
 53            if node._meta is not None:
 54                payload[META] = node._meta
 55            if node.args:
 56                for k, vs in reversed(node.args.items()):
 57                    if type(vs) is list:
 58                        for v in reversed(vs):
 59                            stack.append((v, i, k, True))
 60                    elif vs is not None:
 61                        stack.append((vs, i, k, False))
 62        elif type(node) is exp.DType:
 63            payload[CLASS] = DATA_TYPE
 64            payload[VALUE] = node.value
 65        else:
 66            payload[VALUE] = node
 67
 68        i += 1
 69
 70    return payloads
 71
 72
 73def load(
 74    payloads: list[dict[str, t.Any]] | None,
 75) -> exp.Expr | exp.DType | None:
 76    """
 77    Load a list of dicts generated by dump into an Expr.
 78    """
 79
 80    if not payloads:
 81        return None
 82
 83    payload, *tail = payloads
 84    root = _load(payload)
 85    nodes: list[object] = [root]
 86    for payload in tail:
 87        if CLASS in payload:
 88            node: object = _load(payload)
 89        else:
 90            node = payload[VALUE]
 91
 92        nodes.append(node)
 93        parent = nodes[payload[INDEX]]
 94        arg_key = payload[ARG_KEY]
 95
 96        if payload.get(IS_ARRAY):
 97            parent.append(arg_key, node)
 98        else:
 99            parent.set(arg_key, node)
100
101    return root
102
103
104def _load(payload: dict[str, t.Any]) -> exp.Expr | exp.DType:
105    class_name = payload[CLASS]
106
107    if class_name == DATA_TYPE:
108        return exp.DType(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