sqlglot.typing

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