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]
 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 = {
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"] = node.type.sql()
33        if node.comments:
34            obj["comments"] = node.comments
35        return obj
36    return node
37
38
39def load(obj: JSON) -> Node:
40    """
41    Recursively load a dict (as returned by `dump`) into an AST.
42    """
43    if isinstance(obj, list):
44        return [load(i) for i in obj]
45    if isinstance(obj, dict):
46        class_name = obj["class"]
47
48        if class_name == "DataType.Type":
49            return exp.DataType.Type(obj["value"])
50
51        if "." in class_name:
52            module_path, class_name = class_name.rsplit(".", maxsplit=1)
53            module = __import__(module_path, fromlist=[class_name])
54        else:
55            module = exp
56
57        klass = getattr(module, class_name)
58
59        expression = klass(**{k: load(v) for k, v in obj["args"].items()})
60        type_ = obj.get("type")
61        if type_:
62            expression.type = exp.DataType.build(type_)
63        comments = obj.get("comments")
64        if comments:
65            expression.comments = load(comments)
66        return expression
67    return obj