sqlglot.typing

  1import typing as t
  2
  3from sqlglot import exp
  4from sqlglot.helper import subclasses
  5
  6ExpressionMetadataType = t.Dict[type[exp.Expression], t.Dict[str, t.Any]]
  7
  8TIMESTAMP_EXPRESSIONS = {
  9    exp.CurrentTimestamp,
 10    exp.StrToTime,
 11    exp.TimeStrToTime,
 12    exp.TimestampAdd,
 13    exp.TimestampSub,
 14    exp.UnixToTime,
 15}
 16
 17EXPRESSION_METADATA: ExpressionMetadataType = {
 18    **{
 19        expr_type: {"annotator": lambda self, e: self._annotate_binary(e)}
 20        for expr_type in subclasses(exp.__name__, exp.Binary)
 21    },
 22    **{
 23        expr_type: {"annotator": lambda self, e: self._annotate_unary(e)}
 24        for expr_type in subclasses(exp.__name__, (exp.Unary, exp.Alias))
 25    },
 26    **{
 27        expr_type: {"returns": exp.DataType.Type.BIGINT}
 28        for expr_type in {
 29            exp.ApproxDistinct,
 30            exp.ArraySize,
 31            exp.CountIf,
 32            exp.Int64,
 33            exp.UnixDate,
 34            exp.UnixSeconds,
 35            exp.UnixMicros,
 36            exp.UnixMillis,
 37        }
 38    },
 39    **{
 40        expr_type: {"returns": exp.DataType.Type.BINARY}
 41        for expr_type in {
 42            exp.FromBase32,
 43            exp.FromBase64,
 44        }
 45    },
 46    **{
 47        expr_type: {"returns": exp.DataType.Type.BOOLEAN}
 48        for expr_type in {
 49            exp.All,
 50            exp.Any,
 51            exp.Between,
 52            exp.Boolean,
 53            exp.Contains,
 54            exp.EndsWith,
 55            exp.Exists,
 56            exp.In,
 57            exp.IsInf,
 58            exp.IsNan,
 59            exp.LogicalAnd,
 60            exp.LogicalOr,
 61            exp.RegexpLike,
 62            exp.StartsWith,
 63        }
 64    },
 65    **{
 66        expr_type: {"returns": exp.DataType.Type.DATE}
 67        for expr_type in {
 68            exp.CurrentDate,
 69            exp.Date,
 70            exp.DateFromParts,
 71            exp.DateStrToDate,
 72            exp.DiToDate,
 73            exp.LastDay,
 74            exp.StrToDate,
 75            exp.TimeStrToDate,
 76            exp.TsOrDsToDate,
 77        }
 78    },
 79    **{
 80        expr_type: {"returns": exp.DataType.Type.DATETIME}
 81        for expr_type in {
 82            exp.CurrentDatetime,
 83            exp.Datetime,
 84            exp.DatetimeAdd,
 85            exp.DatetimeSub,
 86        }
 87    },
 88    **{
 89        expr_type: {"returns": exp.DataType.Type.DOUBLE}
 90        for expr_type in {
 91            exp.ApproxQuantile,
 92            exp.Avg,
 93            exp.Cbrt,
 94            exp.Exp,
 95            exp.Kurtosis,
 96            exp.Ln,
 97            exp.Log,
 98            exp.Pi,
 99            exp.Pow,
100            exp.Quantile,
101            exp.Radians,
102            exp.Round,
103            exp.SafeDivide,
104            exp.Sqrt,
105            exp.Stddev,
106            exp.StddevPop,
107            exp.StddevSamp,
108            exp.ToDouble,
109            exp.Variance,
110            exp.VariancePop,
111            exp.Skewness,
112        }
113    },
114    **{
115        expr_type: {"returns": exp.DataType.Type.INT}
116        for expr_type in {
117            exp.Ascii,
118            exp.BitLength,
119            exp.Ceil,
120            exp.DatetimeDiff,
121            exp.Getbit,
122            exp.TimestampDiff,
123            exp.TimeDiff,
124            exp.Unicode,
125            exp.DateToDi,
126            exp.Levenshtein,
127            exp.Length,
128            exp.Sign,
129            exp.StrPosition,
130            exp.TsOrDiToDi,
131            exp.Quarter,
132        }
133    },
134    **{
135        expr_type: {"returns": exp.DataType.Type.INTERVAL}
136        for expr_type in {
137            exp.Interval,
138            exp.JustifyDays,
139            exp.JustifyHours,
140            exp.JustifyInterval,
141            exp.MakeInterval,
142        }
143    },
144    **{
145        expr_type: {"returns": exp.DataType.Type.JSON}
146        for expr_type in {
147            exp.ParseJSON,
148        }
149    },
150    **{
151        expr_type: {"returns": exp.DataType.Type.TIME}
152        for expr_type in {
153            exp.CurrentTime,
154            exp.Localtime,
155            exp.Time,
156            exp.TimeAdd,
157            exp.TimeSub,
158        }
159    },
160    **{
161        expr_type: {"returns": exp.DataType.Type.TIMESTAMPLTZ}
162        for expr_type in {
163            exp.TimestampLtzFromParts,
164        }
165    },
166    **{
167        expr_type: {"returns": exp.DataType.Type.TIMESTAMPTZ}
168        for expr_type in {
169            exp.CurrentTimestampLTZ,
170            exp.TimestampTzFromParts,
171        }
172    },
173    **{expr_type: {"returns": exp.DataType.Type.TIMESTAMP} for expr_type in TIMESTAMP_EXPRESSIONS},
174    **{
175        expr_type: {"returns": exp.DataType.Type.TINYINT}
176        for expr_type in {
177            exp.Day,
178            exp.DayOfMonth,
179            exp.DayOfWeek,
180            exp.DayOfWeekIso,
181            exp.DayOfYear,
182            exp.Month,
183            exp.Week,
184            exp.WeekOfYear,
185            exp.Year,
186            exp.YearOfWeek,
187            exp.YearOfWeekIso,
188        }
189    },
190    **{
191        expr_type: {"returns": exp.DataType.Type.VARCHAR}
192        for expr_type in {
193            exp.ArrayToString,
194            exp.Concat,
195            exp.ConcatWs,
196            exp.Chr,
197            exp.Dayname,
198            exp.DateToDateStr,
199            exp.DPipe,
200            exp.GroupConcat,
201            exp.Initcap,
202            exp.Lower,
203            exp.SHA,
204            exp.SHA2,
205            exp.Substring,
206            exp.String,
207            exp.TimeToStr,
208            exp.TimeToTimeStr,
209            exp.Trim,
210            exp.ToBase32,
211            exp.ToBase64,
212            exp.TsOrDsToDateStr,
213            exp.UnixToStr,
214            exp.UnixToTimeStr,
215            exp.Upper,
216            exp.RawString,
217            exp.Space,
218        }
219    },
220    **{
221        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this")}
222        for expr_type in {
223            exp.Abs,
224            exp.AnyValue,
225            exp.ArrayConcatAgg,
226            exp.ArrayReverse,
227            exp.ArraySlice,
228            exp.Filter,
229            exp.HavingMax,
230            exp.LastValue,
231            exp.Limit,
232            exp.Order,
233            exp.SortArray,
234            exp.Window,
235        }
236    },
237    **{
238        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this", "expressions")}
239        for expr_type in {
240            exp.ArrayConcat,
241            exp.Coalesce,
242            exp.Greatest,
243            exp.Least,
244            exp.Max,
245            exp.Min,
246        }
247    },
248    **{
249        expr_type: {"annotator": lambda self, e: self._annotate_by_array_element(e)}
250        for expr_type in {
251            exp.ArrayFirst,
252            exp.ArrayLast,
253        }
254    },
255    exp.Anonymous: {"annotator": lambda self, e: self._set_type(e, self.schema.get_udf_type(e))},
256    **{
257        expr_type: {"annotator": lambda self, e: self._annotate_timeunit(e)}
258        for expr_type in {
259            exp.DateAdd,
260            exp.DateSub,
261            exp.DateTrunc,
262        }
263    },
264    **{
265        expr_type: {"annotator": lambda self, e: self._set_type(e, e.args["to"])}
266        for expr_type in {
267            exp.Cast,
268            exp.TryCast,
269        }
270    },
271    **{
272        expr_type: {"annotator": lambda self, e: self._annotate_map(e)}
273        for expr_type in {
274            exp.Map,
275            exp.VarMap,
276        }
277    },
278    exp.Array: {"annotator": lambda self, e: self._annotate_by_args(e, "expressions", array=True)},
279    exp.ArrayAgg: {"annotator": lambda self, e: self._annotate_by_args(e, "this", array=True)},
280    exp.Bracket: {"annotator": lambda self, e: self._annotate_bracket(e)},
281    exp.Case: {
282        "annotator": lambda self, e: self._annotate_by_args(
283            e, *[if_expr.args["true"] for if_expr in e.args["ifs"]], "default"
284        )
285    },
286    exp.Count: {
287        "annotator": lambda self, e: self._set_type(
288            e, exp.DataType.Type.BIGINT if e.args.get("big_int") else exp.DataType.Type.INT
289        )
290    },
291    exp.DateDiff: {
292        "annotator": lambda self, e: self._set_type(
293            e, exp.DataType.Type.BIGINT if e.args.get("big_int") else exp.DataType.Type.INT
294        )
295    },
296    exp.DataType: {"annotator": lambda self, e: self._set_type(e, e.copy())},
297    exp.Div: {"annotator": lambda self, e: self._annotate_div(e)},
298    exp.Distinct: {"annotator": lambda self, e: self._annotate_by_args(e, "expressions")},
299    exp.Dot: {"annotator": lambda self, e: self._annotate_dot(e)},
300    exp.Explode: {"annotator": lambda self, e: self._annotate_explode(e)},
301    exp.Extract: {"annotator": lambda self, e: self._annotate_extract(e)},
302    exp.HexString: {
303        "annotator": lambda self, e: self._set_type(
304            e,
305            exp.DataType.Type.BIGINT if e.args.get("is_integer") else exp.DataType.Type.BINARY,
306        )
307    },
308    exp.GenerateSeries: {
309        "annotator": lambda self, e: self._annotate_by_args(e, "start", "end", "step", array=True)
310    },
311    exp.GenerateDateArray: {
312        "annotator": lambda self, e: self._set_type(e, exp.DataType.build("ARRAY<DATE>"))
313    },
314    exp.GenerateTimestampArray: {
315        "annotator": lambda self, e: self._set_type(e, exp.DataType.build("ARRAY<TIMESTAMP>"))
316    },
317    exp.If: {"annotator": lambda self, e: self._annotate_by_args(e, "true", "false")},
318    exp.Literal: {"annotator": lambda self, e: self._annotate_literal(e)},
319    exp.Null: {"returns": exp.DataType.Type.NULL},
320    exp.Nullif: {"annotator": lambda self, e: self._annotate_by_args(e, "this", "expression")},
321    exp.PropertyEQ: {"annotator": lambda self, e: self._annotate_by_args(e, "expression")},
322    exp.Struct: {"annotator": lambda self, e: self._annotate_struct(e)},
323    exp.Sum: {
324        "annotator": lambda self, e: self._annotate_by_args(e, "this", "expressions", promote=True)
325    },
326    exp.Timestamp: {
327        "annotator": lambda self, e: self._set_type(
328            e,
329            exp.DataType.Type.TIMESTAMPTZ if e.args.get("with_tz") else exp.DataType.Type.TIMESTAMP,
330        )
331    },
332    exp.ToMap: {"annotator": lambda self, e: self._annotate_to_map(e)},
333    exp.Unnest: {"annotator": lambda self, e: self._annotate_unnest(e)},
334    exp.Subquery: {"annotator": lambda self, e: self._annotate_subquery(e)},
335}