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