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