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