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