1"""sqlglot expressions datatypes."""
  2
  3from __future__ import annotations
  4
  5import typing as t
  6from enum import auto
  7
  8from sqlglot.helper import AutoName
  9from sqlglot.errors import ErrorLevel, ParseError
 10from sqlglot.expressions.core import (
 11    Expr,
 12    Expression,
 13    _TimeUnit,
 14    Identifier,
 15    Dot,
 16    maybe_copy,
 17)
 18from builtins import type as Type
 19
 20if t.TYPE_CHECKING:
 21    from sqlglot._typing import DataTypeArgs
 22    from sqlglot.dialects.dialect import DialectType
 23    from typing_extensions import Self, Unpack
 24
 25
 26class DataTypeParam(Expression):
 27    arg_types = {"this": True, "expression": False}
 28
 29    @property
 30    def this(self) -> Expr:
 31        return self.args["this"]
 32
 33    @property
 34    def name(self) -> str:
 35        return self.this.name
 36
 37
 38class DType(AutoName):
 39    ARRAY = auto()
 40    AGGREGATEFUNCTION = auto()
 41    SIMPLEAGGREGATEFUNCTION = auto()
 42    BIGDECIMAL = auto()
 43    BIGINT = auto()
 44    BIGNUM = auto()
 45    BIGSERIAL = auto()
 46    BINARY = auto()
 47    BIT = auto()
 48    BLOB = auto()
 49    BOOLEAN = auto()
 50    BPCHAR = auto()
 51    CHAR = auto()
 52    CHARACTER_SET = auto()
 53    DATE = auto()
 54    DATE32 = auto()
 55    DATEMULTIRANGE = auto()
 56    DATERANGE = auto()
 57    DATETIME = auto()
 58    DATETIME2 = auto()
 59    DATETIME64 = auto()
 60    DECIMAL = auto()
 61    DECIMAL32 = auto()
 62    DECIMAL64 = auto()
 63    DECIMAL128 = auto()
 64    DECIMAL256 = auto()
 65    DECFLOAT = auto()
 66    DOUBLE = auto()
 67    DYNAMIC = auto()
 68    ENUM = auto()
 69    ENUM8 = auto()
 70    ENUM16 = auto()
 71    FILE = auto()
 72    FIXEDSTRING = auto()
 73    FLOAT = auto()
 74    GEOGRAPHY = auto()
 75    GEOGRAPHYPOINT = auto()
 76    GEOMETRY = auto()
 77    POINT = auto()
 78    RING = auto()
 79    LINESTRING = auto()
 80    MULTILINESTRING = auto()
 81    POLYGON = auto()
 82    MULTIPOLYGON = auto()
 83    HLLSKETCH = auto()
 84    HSTORE = auto()
 85    IMAGE = auto()
 86    INET = auto()
 87    INT = auto()
 88    INT128 = auto()
 89    INT256 = auto()
 90    INT4MULTIRANGE = auto()
 91    INT4RANGE = auto()
 92    INT8MULTIRANGE = auto()
 93    INT8RANGE = auto()
 94    INTERVAL = auto()
 95    IPADDRESS = auto()
 96    IPPREFIX = auto()
 97    IPV4 = auto()
 98    IPV6 = auto()
 99    JSON = auto()
100    JSONB = auto()
101    LIST = auto()
102    LONGBLOB = auto()
103    LONGTEXT = auto()
104    LOWCARDINALITY = auto()
105    MAP = auto()
106    MEDIUMBLOB = auto()
107    MEDIUMINT = auto()
108    MEDIUMTEXT = auto()
109    MONEY = auto()
110    NAME = auto()
111    NCHAR = auto()
112    NESTED = auto()
113    NOTHING = auto()
114    NULL = auto()
115    NUMMULTIRANGE = auto()
116    NUMRANGE = auto()
117    NVARCHAR = auto()
118    OBJECT = auto()
119    RANGE = auto()
120    REF = auto()
121    ROWVERSION = auto()
122    SERIAL = auto()
123    SET = auto()
124    SMALLDATETIME = auto()
125    SMALLINT = auto()
126    SMALLMONEY = auto()
127    SMALLSERIAL = auto()
128    STRUCT = auto()
129    SUPER = auto()
130    TEXT = auto()
131    TINYBLOB = auto()
132    TINYTEXT = auto()
133    TIME = auto()
134    TIMETZ = auto()
135    TIME_NS = auto()
136    TIMESTAMP = auto()
137    TIMESTAMPNTZ = auto()
138    TIMESTAMPLTZ = auto()
139    TIMESTAMPTZ = auto()
140    TIMESTAMP_S = auto()
141    TIMESTAMP_MS = auto()
142    TIMESTAMP_NS = auto()
143    TINYINT = auto()
144    TSMULTIRANGE = auto()
145    TSRANGE = auto()
146    TSTZMULTIRANGE = auto()
147    TSTZRANGE = auto()
148    UBIGINT = auto()
149    UINT = auto()
150    UINT128 = auto()
151    UINT256 = auto()
152    UMEDIUMINT = auto()
153    UDECIMAL = auto()
154    UDOUBLE = auto()
155    UNION = auto()
156    UNKNOWN = auto()  # Sentinel value, useful for type annotation
157    USERDEFINED = "USER-DEFINED"
158    USMALLINT = auto()
159    UTINYINT = auto()
160    UUID = auto()
161    VARBINARY = auto()
162    VARCHAR = auto()
163    VARIANT = auto()
164    VECTOR = auto()
165    XML = auto()
166    YEAR = auto()
167    TDIGEST = auto()
168
169    def into_expr(self, **kwargs: object) -> DataType:
170        """Converts this `DType` into a `DataType` instance.
171
172        Args:
173            **kwargs (object): additional arguments to pass in the constructor of DataType.
174        Returns:
175            DataType: the resulting `DataType` instance.
176        """
177        return DataType(this=self).set_kwargs(kwargs)
178
179
180class DataType(Expression):
181    arg_types = {
182        "this": True,
183        "expressions": False,
184        "nested": False,
185        "values": False,
186        "kind": False,
187        "nullable": False,
188        "collate": False,
189    }
190
191    is_data_type: t.ClassVar[bool] = True
192
193    Type: t.ClassVar[Type[DType]] = DType
194
195    STRUCT_TYPES: t.ClassVar[set[DType]] = {
196        DType.FILE,
197        DType.NESTED,
198        DType.OBJECT,
199        DType.STRUCT,
200        DType.UNION,
201    }
202
203    ARRAY_TYPES: t.ClassVar[set[DType]] = {
204        DType.ARRAY,
205        DType.LIST,
206    }
207
208    NESTED_TYPES: t.ClassVar[set[DType]] = {
209        DType.FILE,
210        DType.NESTED,
211        DType.OBJECT,
212        DType.STRUCT,
213        DType.UNION,
214        DType.ARRAY,
215        DType.LIST,
216        DType.MAP,
217    }
218
219    TEXT_TYPES: t.ClassVar[set[DType]] = {
220        DType.CHAR,
221        DType.NCHAR,
222        DType.NVARCHAR,
223        DType.TEXT,
224        DType.TINYTEXT,
225        DType.MEDIUMTEXT,
226        DType.LONGTEXT,
227        DType.VARCHAR,
228        DType.NAME,
229    }
230
231    BINARY_TYPES: t.ClassVar[set[DType]] = {
232        DType.BINARY,
233        DType.VARBINARY,
234        DType.TINYBLOB,
235        DType.BLOB,
236        DType.MEDIUMBLOB,
237        DType.LONGBLOB,
238    }
239
240    SIGNED_INTEGER_TYPES: t.ClassVar[set[DType]] = {
241        DType.BIGINT,
242        DType.INT,
243        DType.INT128,
244        DType.INT256,
245        DType.MEDIUMINT,
246        DType.SMALLINT,
247        DType.TINYINT,
248    }
249
250    UNSIGNED_INTEGER_TYPES: t.ClassVar[set[DType]] = {
251        DType.UBIGINT,
252        DType.UINT,
253        DType.UINT128,
254        DType.UINT256,
255        DType.UMEDIUMINT,
256        DType.USMALLINT,
257        DType.UTINYINT,
258    }
259
260    INTEGER_TYPES: t.ClassVar[set[DType]] = {
261        DType.BIGINT,
262        DType.INT,
263        DType.INT128,
264        DType.INT256,
265        DType.MEDIUMINT,
266        DType.SMALLINT,
267        DType.TINYINT,
268        DType.UBIGINT,
269        DType.UINT,
270        DType.UINT128,
271        DType.UINT256,
272        DType.UMEDIUMINT,
273        DType.USMALLINT,
274        DType.UTINYINT,
275        DType.BIT,
276    }
277
278    FLOAT_TYPES: t.ClassVar[set[DType]] = {
279        DType.DOUBLE,
280        DType.FLOAT,
281    }
282
283    REAL_TYPES: t.ClassVar[set[DType]] = {
284        DType.DOUBLE,
285        DType.FLOAT,
286        DType.BIGDECIMAL,
287        DType.DECIMAL,
288        DType.DECIMAL32,
289        DType.DECIMAL64,
290        DType.DECIMAL128,
291        DType.DECIMAL256,
292        DType.DECFLOAT,
293        DType.MONEY,
294        DType.SMALLMONEY,
295        DType.UDECIMAL,
296        DType.UDOUBLE,
297    }
298
299    NUMERIC_TYPES: t.ClassVar[set[DType]] = {
300        DType.BIGINT,
301        DType.INT,
302        DType.INT128,
303        DType.INT256,
304        DType.MEDIUMINT,
305        DType.SMALLINT,
306        DType.TINYINT,
307        DType.UBIGINT,
308        DType.UINT,
309        DType.UINT128,
310        DType.UINT256,
311        DType.UMEDIUMINT,
312        DType.USMALLINT,
313        DType.UTINYINT,
314        DType.BIT,
315        DType.DOUBLE,
316        DType.FLOAT,
317        DType.BIGDECIMAL,
318        DType.DECIMAL,
319        DType.DECIMAL32,
320        DType.DECIMAL64,
321        DType.DECIMAL128,
322        DType.DECIMAL256,
323        DType.DECFLOAT,
324        DType.MONEY,
325        DType.SMALLMONEY,
326        DType.UDECIMAL,
327        DType.UDOUBLE,
328    }
329
330    TEMPORAL_TYPES: t.ClassVar[set[DType]] = {
331        DType.DATE,
332        DType.DATE32,
333        DType.DATETIME,
334        DType.DATETIME2,
335        DType.DATETIME64,
336        DType.SMALLDATETIME,
337        DType.TIME,
338        DType.TIMESTAMP,
339        DType.TIMESTAMPNTZ,
340        DType.TIMESTAMPLTZ,
341        DType.TIMESTAMPTZ,
342        DType.TIMESTAMP_MS,
343        DType.TIMESTAMP_NS,
344        DType.TIMESTAMP_S,
345        DType.TIMETZ,
346    }
347
348    @classmethod
349    def build(
350        cls,
351        dtype: DATA_TYPE,
352        dialect: DialectType = None,
353        udt: bool = False,
354        copy: bool = True,
355        **kwargs: Unpack[DataTypeArgs],
356    ) -> Self:
357        """
358        Constructs a DataType object.
359
360        Args:
361            dtype: the data type of interest.
362            dialect: the dialect to use for parsing `dtype`, in case it's a string.
363            udt: when set to True, `dtype` will be used as-is if it can't be parsed into a
364                DataType, thus creating a user-defined type.
365            copy: whether to copy the data type.
366            kwargs: additional arguments to pass in the constructor of DataType.
367
368        Returns:
369            The constructed DataType object.
370        """
371        if isinstance(dtype, str):
372            return cls.from_str(dtype, dialect, udt, **kwargs)
373        elif isinstance(dtype, DType):
374            data_type_exp = cls(this=dtype)
375            if kwargs:
376                for k, v in kwargs.items():
377                    data_type_exp.set(k, v)
378            return data_type_exp
379        elif isinstance(dtype, (Identifier, Dot)) and udt:
380            return cls(this=DType.USERDEFINED, kind=dtype, **kwargs)
381        elif isinstance(dtype, cls):
382            return maybe_copy(dtype, copy)
383        else:
384            raise ValueError(f"Invalid data type: {type(dtype)}. Expected str or DType")
385
386    @classmethod
387    def from_str(
388        cls,
389        dtype: str,
390        dialect: DialectType = None,
391        udt: bool = False,
392        **kwargs: Unpack[DataTypeArgs],
393    ) -> Self:
394        """
395        Constructs a `DataType` object from a `str` representation.
396
397        Args:
398            dtype: the data type of interest.
399            dialect: the dialect to use for parsing `dtype`.
400            udt: when set to True, `dtype` will be used as-is if it can't be parsed into a
401                `DataType`, thus creating a user-defined type.
402            kwargs: additional arguments to pass in the constructor of `DataType`.
403
404        Returns:
405            The constructed `DataType` object.
406        """
407        from sqlglot import parse_one
408
409        if dtype.upper() == "UNKNOWN":
410            return cls(this=DType.UNKNOWN, **kwargs)
411        try:
412            return parse_one(
413                dtype, read=dialect, into=cls, error_level=ErrorLevel.IGNORE
414            ).set_kwargs(kwargs)
415        except ParseError:
416            if udt:
417                return cls(this=DType.USERDEFINED, kind=dtype, **kwargs)
418            raise
419
420    def is_type(self, *dtypes: DATA_TYPE, check_nullable: bool = False) -> bool:
421        """
422        Checks whether this DataType matches one of the provided data types. Nested types or precision
423        will be compared using "structural equivalence" semantics, so e.g. array<int> != array<float>.
424
425        Args:
426            dtypes: the data types to compare this DataType to.
427            check_nullable: whether to take the NULLABLE type constructor into account for the comparison.
428                If false, it means that NULLABLE<INT> is equivalent to INT.
429
430        Returns:
431            True, if and only if there is a type in `dtypes` which is equal to this DataType.
432        """
433        self_is_nullable: bool | None = self.args.get("nullable")
434        for dtype in dtypes:
435            other_type = DataType.build(dtype, copy=False, udt=True)
436            other_is_nullable: bool | None = other_type.args.get("nullable")
437            if (
438                other_type.expressions
439                or (check_nullable and (self_is_nullable or other_is_nullable))
440                or self.this == DType.USERDEFINED
441                or other_type.this == DType.USERDEFINED
442            ):
443                matches = self == other_type
444            else:
445                matches = self.this == other_type.this
446
447            if matches:
448                return True
449        return False
450
451
452class PseudoType(DataType):
453    arg_types = {"this": True}
454
455
456class ObjectIdentifier(DataType):
457    arg_types = {"this": True}
458
459
460class IntervalSpan(DataType):
461    arg_types = {"this": True, "expression": True}
462
463
464class Interval(_TimeUnit):
465    # fraction_in_string: a fractional amount was written with its unit in one string, e.g.
466    # INTERVAL '1.5 DAY', which differs from INTERVAL '1.5' DAY in e.g. DuckDB and Postgres
467    arg_types = {"this": False, "unit": False, "fraction_in_string": False}
468
469
470DATA_TYPE = t.Union[str, Identifier, Dot, DataType, DType]