sqlglot.optimizer.canonicalize

 1import itertools
 2
 3from sqlglot import exp
 4
 5
 6def canonicalize(expression: exp.Expression) -> exp.Expression:
 7    """Converts a sql expression into a standard form.
 8
 9    This method relies on annotate_types because many of the
10    conversions rely on type inference.
11
12    Args:
13        expression: The expression to canonicalize.
14    """
15    exp.replace_children(expression, canonicalize)
16
17    expression = add_text_to_concat(expression)
18    expression = coerce_type(expression)
19    expression = remove_redundant_casts(expression)
20
21    if isinstance(expression, exp.Identifier):
22        expression.set("quoted", True)
23
24    return expression
25
26
27def add_text_to_concat(node: exp.Expression) -> exp.Expression:
28    if isinstance(node, exp.Add) and node.type and node.type.this in exp.DataType.TEXT_TYPES:
29        node = exp.Concat(this=node.this, expression=node.expression)
30    return node
31
32
33def coerce_type(node: exp.Expression) -> exp.Expression:
34    if isinstance(node, exp.Binary):
35        _coerce_date(node.left, node.right)
36    elif isinstance(node, exp.Between):
37        _coerce_date(node.this, node.args["low"])
38    elif isinstance(node, exp.Extract):
39        if node.expression.type.this not in exp.DataType.TEMPORAL_TYPES:
40            _replace_cast(node.expression, "datetime")
41    return node
42
43
44def remove_redundant_casts(expression: exp.Expression) -> exp.Expression:
45    if (
46        isinstance(expression, exp.Cast)
47        and expression.to.type
48        and expression.this.type
49        and expression.to.type.this == expression.this.type.this
50    ):
51        return expression.this
52    return expression
53
54
55def _coerce_date(a: exp.Expression, b: exp.Expression) -> None:
56    for a, b in itertools.permutations([a, b]):
57        if (
58            a.type
59            and a.type.this == exp.DataType.Type.DATE
60            and b.type
61            and b.type.this != exp.DataType.Type.DATE
62        ):
63            _replace_cast(b, "date")
64
65
66def _replace_cast(node: exp.Expression, to: str) -> None:
67    data_type = exp.DataType.build(to)
68    cast = exp.Cast(this=node.copy(), to=data_type)
69    cast.type = data_type
70    node.replace(cast)