sqlglot.serde

 1from __future__ import annotations
 2
 3import typing as t
 4
 5from sqlglot import expressions as exp
 6
 7if t.TYPE_CHECKING:
 8    JSON = t.Union[dict, list, str, float, int, bool, None]
 9    Node = t.Union[t.List["Node"], exp.DataType.Type, exp.Expression, JSON]
10
11
12def dump(node: Node) -> JSON:
13    """
14    Recursively dump an AST into a JSON-serializable dict.
15    """
16    if isinstance(node, list):
17        return [dump(i) for i in node]
18    if isinstance(node, exp.DataType.Type):
19        return {
20            "class": "DataType.Type",
21            "value": node.value,
22        }
23    if isinstance(node, exp.Expression):
24        klass = node.__class__.__qualname__
25        if node.__class__.__module__ != exp.__name__:
26            klass = f"{node.__module__}.{klass}"
27        obj: t.Dict = {
28            "class": klass,
29            "args": {k: dump(v) for k, v in node.args.items() if v is not None and v != []},
30        }
31        if node.type:
32            obj["type"] = dump(node.type)
33        if node.comments:
34            obj["comments"] = node.comments
35        if node._meta is not None:
36            obj["meta"] = node._meta
37
38        return obj
39    return node
40
41
42def load(obj: JSON) -> Node:
43    """
44    Recursively load a dict (as returned by `dump`) into an AST.
45    """
46    if isinstance(obj, list):
47        return [load(i) for i in obj]
48    if isinstance(obj, dict):
49        class_name = obj["class"]
50
51        if class_name == "DataType.Type":
52            return exp.DataType.Type(obj["value"])
53
54        if "." in class_name:
55            module_path, class_name = class_name.rsplit(".", maxsplit=1)
56            module = __import__(module_path, fromlist=[class_name])
57        else:
58            module = exp
59
60        klass = getattr(module, class_name)
61
62        expression = klass(**{k: load(v) for k, v in obj["args"].items()})
63        expression.type = t.cast(exp.DataType, load(obj.get("type")))
64        expression.comments = obj.get("comments")
65        expression._meta = obj.get("meta")
66
67        return expression
68    return obj