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