sqlglot.optimizer.canonicalize

 1from __future__ import annotations
 2
 3import itertools
 4
 5from sqlglot import exp
 6
 7
 8def canonicalize(expression: exp.Expression) -> exp.Expression:
 9    """Converts a sql expression into a standard form.
10
11    This method relies on annotate_types because many of the
12    conversions rely on type inference.
13
14    Args:
15        expression: The expression to canonicalize.
16    """
17    exp.replace_children(expression, canonicalize)
18
19    expression = add_text_to_concat(expression)
20    expression = coerce_type(expression)
21    expression = remove_redundant_casts(expression)
22    expression = ensure_bool_predicates(expression)
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(expressions=[node.left, node.right])
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 ensure_bool_predicates(expression: exp.Expression) -> exp.Expression:
56    if isinstance(expression, exp.Connector):
57        _replace_int_predicate(expression.left)
58        _replace_int_predicate(expression.right)
59
60    elif isinstance(expression, (exp.Where, exp.Having)):
61        _replace_int_predicate(expression.this)
62
63    return expression
64
65
66def _coerce_date(a: exp.Expression, b: exp.Expression) -> None:
67    for a, b in itertools.permutations([a, b]):
68        if (
69            a.type
70            and a.type.this == exp.DataType.Type.DATE
71            and b.type
72            and b.type.this not in (exp.DataType.Type.DATE, exp.DataType.Type.INTERVAL)
73        ):
74            _replace_cast(b, "date")
75
76
77def _replace_cast(node: exp.Expression, to: str) -> None:
78    data_type = exp.DataType.build(to)
79    cast = exp.Cast(this=node.copy(), to=data_type)
80    cast.type = data_type
81    node.replace(cast)
82
83
84def _replace_int_predicate(expression: exp.Expression) -> None:
85    if expression.type and expression.type.this in exp.DataType.INTEGER_TYPES:
86        expression.replace(exp.NEQ(this=expression.copy(), expression=exp.Literal.number(0)))