sqlglot.helper

  1from __future__ import annotations
  2
  3import datetime
  4import inspect
  5import logging
  6import re
  7import sys
  8import typing as t
  9from collections.abc import Collection
 10from contextlib import contextmanager
 11from copy import copy
 12from enum import Enum
 13from itertools import count
 14
 15if t.TYPE_CHECKING:
 16    from sqlglot import exp
 17    from sqlglot._typing import A, E, T
 18    from sqlglot.expressions import Expression
 19
 20
 21CAMEL_CASE_PATTERN = re.compile("(?<!^)(?=[A-Z])")
 22PYTHON_VERSION = sys.version_info[:2]
 23logger = logging.getLogger("sqlglot")
 24
 25
 26class AutoName(Enum):
 27    """
 28    This is used for creating Enum classes where `auto()` is the string form
 29    of the corresponding enum's identifier (e.g. FOO.value results in "FOO").
 30
 31    Reference: https://docs.python.org/3/howto/enum.html#using-automatic-values
 32    """
 33
 34    def _generate_next_value_(name, _start, _count, _last_values):
 35        return name
 36
 37
 38class classproperty(property):
 39    """
 40    Similar to a normal property but works for class methods
 41    """
 42
 43    def __get__(self, obj: t.Any, owner: t.Any = None) -> t.Any:
 44        return classmethod(self.fget).__get__(None, owner)()  # type: ignore
 45
 46
 47def seq_get(seq: t.Sequence[T], index: int) -> t.Optional[T]:
 48    """Returns the value in `seq` at position `index`, or `None` if `index` is out of bounds."""
 49    try:
 50        return seq[index]
 51    except IndexError:
 52        return None
 53
 54
 55@t.overload
 56def ensure_list(value: t.Collection[T]) -> t.List[T]:
 57    ...
 58
 59
 60@t.overload
 61def ensure_list(value: T) -> t.List[T]:
 62    ...
 63
 64
 65def ensure_list(value):
 66    """
 67    Ensures that a value is a list, otherwise casts or wraps it into one.
 68
 69    Args:
 70        value: The value of interest.
 71
 72    Returns:
 73        The value cast as a list if it's a list or a tuple, or else the value wrapped in a list.
 74    """
 75    if value is None:
 76        return []
 77    if isinstance(value, (list, tuple)):
 78        return list(value)
 79
 80    return [value]
 81
 82
 83@t.overload
 84def ensure_collection(value: t.Collection[T]) -> t.Collection[T]:
 85    ...
 86
 87
 88@t.overload
 89def ensure_collection(value: T) -> t.Collection[T]:
 90    ...
 91
 92
 93def ensure_collection(value):
 94    """
 95    Ensures that a value is a collection (excluding `str` and `bytes`), otherwise wraps it into a list.
 96
 97    Args:
 98        value: The value of interest.
 99
100    Returns:
101        The value if it's a collection, or else the value wrapped in a list.
102    """
103    if value is None:
104        return []
105    return (
106        value if isinstance(value, Collection) and not isinstance(value, (str, bytes)) else [value]
107    )
108
109
110def csv(*args: str, sep: str = ", ") -> str:
111    """
112    Formats any number of string arguments as CSV.
113
114    Args:
115        args: The string arguments to format.
116        sep: The argument separator.
117
118    Returns:
119        The arguments formatted as a CSV string.
120    """
121    return sep.join(arg for arg in args if arg)
122
123
124def subclasses(
125    module_name: str,
126    classes: t.Type | t.Tuple[t.Type, ...],
127    exclude: t.Type | t.Tuple[t.Type, ...] = (),
128) -> t.List[t.Type]:
129    """
130    Returns all subclasses for a collection of classes, possibly excluding some of them.
131
132    Args:
133        module_name: The name of the module to search for subclasses in.
134        classes: Class(es) we want to find the subclasses of.
135        exclude: Class(es) we want to exclude from the returned list.
136
137    Returns:
138        The target subclasses.
139    """
140    return [
141        obj
142        for _, obj in inspect.getmembers(
143            sys.modules[module_name],
144            lambda obj: inspect.isclass(obj) and issubclass(obj, classes) and obj not in exclude,
145        )
146    ]
147
148
149def apply_index_offset(
150    this: exp.Expression,
151    expressions: t.List[E],
152    offset: int,
153) -> t.List[E]:
154    """
155    Applies an offset to a given integer literal expression.
156
157    Args:
158        this: The target of the index.
159        expressions: The expression the offset will be applied to, wrapped in a list.
160        offset: The offset that will be applied.
161
162    Returns:
163        The original expression with the offset applied to it, wrapped in a list. If the provided
164        `expressions` argument contains more than one expression, it's returned unaffected.
165    """
166    if not offset or len(expressions) != 1:
167        return expressions
168
169    expression = expressions[0]
170
171    from sqlglot import exp
172    from sqlglot.optimizer.annotate_types import annotate_types
173    from sqlglot.optimizer.simplify import simplify
174
175    if not this.type:
176        annotate_types(this)
177
178    if t.cast(exp.DataType, this.type).this not in (
179        exp.DataType.Type.UNKNOWN,
180        exp.DataType.Type.ARRAY,
181    ):
182        return expressions
183
184    if not expression.type:
185        annotate_types(expression)
186    if t.cast(exp.DataType, expression.type).this in exp.DataType.INTEGER_TYPES:
187        logger.warning("Applying array index offset (%s)", offset)
188        expression = simplify(exp.Add(this=expression, expression=exp.Literal.number(offset)))
189        return [expression]
190
191    return expressions
192
193
194def camel_to_snake_case(name: str) -> str:
195    """Converts `name` from camelCase to snake_case and returns the result."""
196    return CAMEL_CASE_PATTERN.sub("_", name).upper()
197
198
199def while_changing(expression: Expression, func: t.Callable[[Expression], E]) -> E:
200    """
201    Applies a transformation to a given expression until a fix point is reached.
202
203    Args:
204        expression: The expression to be transformed.
205        func: The transformation to be applied.
206
207    Returns:
208        The transformed expression.
209    """
210    while True:
211        for n, *_ in reversed(tuple(expression.walk())):
212            n._hash = hash(n)
213
214        start = hash(expression)
215        expression = func(expression)
216
217        for n, *_ in expression.walk():
218            n._hash = None
219        if start == hash(expression):
220            break
221
222    return expression
223
224
225def tsort(dag: t.Dict[T, t.Set[T]]) -> t.List[T]:
226    """
227    Sorts a given directed acyclic graph in topological order.
228
229    Args:
230        dag: The graph to be sorted.
231
232    Returns:
233        A list that contains all of the graph's nodes in topological order.
234    """
235    result = []
236
237    for node, deps in tuple(dag.items()):
238        for dep in deps:
239            if not dep in dag:
240                dag[dep] = set()
241
242    while dag:
243        current = {node for node, deps in dag.items() if not deps}
244
245        if not current:
246            raise ValueError("Cycle error")
247
248        for node in current:
249            dag.pop(node)
250
251        for deps in dag.values():
252            deps -= current
253
254        result.extend(sorted(current))  # type: ignore
255
256    return result
257
258
259def open_file(file_name: str) -> t.TextIO:
260    """Open a file that may be compressed as gzip and return it in universal newline mode."""
261    with open(file_name, "rb") as f:
262        gzipped = f.read(2) == b"\x1f\x8b"
263
264    if gzipped:
265        import gzip
266
267        return gzip.open(file_name, "rt", newline="")
268
269    return open(file_name, encoding="utf-8", newline="")
270
271
272@contextmanager
273def csv_reader(read_csv: exp.ReadCSV) -> t.Any:
274    """
275    Returns a csv reader given the expression `READ_CSV(name, ['delimiter', '|', ...])`.
276
277    Args:
278        read_csv: A `ReadCSV` function call.
279
280    Yields:
281        A python csv reader.
282    """
283    args = read_csv.expressions
284    file = open_file(read_csv.name)
285
286    delimiter = ","
287    args = iter(arg.name for arg in args)  # type: ignore
288    for k, v in zip(args, args):
289        if k == "delimiter":
290            delimiter = v
291
292    try:
293        import csv as csv_
294
295        yield csv_.reader(file, delimiter=delimiter)
296    finally:
297        file.close()
298
299
300def find_new_name(taken: t.Collection[str], base: str) -> str:
301    """
302    Searches for a new name.
303
304    Args:
305        taken: A collection of taken names.
306        base: Base name to alter.
307
308    Returns:
309        The new, available name.
310    """
311    if base not in taken:
312        return base
313
314    i = 2
315    new = f"{base}_{i}"
316    while new in taken:
317        i += 1
318        new = f"{base}_{i}"
319
320    return new
321
322
323def name_sequence(prefix: str) -> t.Callable[[], str]:
324    """Returns a name generator given a prefix (e.g. a0, a1, a2, ... if the prefix is "a")."""
325    sequence = count()
326    return lambda: f"{prefix}{next(sequence)}"
327
328
329def object_to_dict(obj: t.Any, **kwargs) -> t.Dict:
330    """Returns a dictionary created from an object's attributes."""
331    return {
332        **{k: v.copy() if hasattr(v, "copy") else copy(v) for k, v in vars(obj).items()},
333        **kwargs,
334    }
335
336
337def split_num_words(
338    value: str, sep: str, min_num_words: int, fill_from_start: bool = True
339) -> t.List[t.Optional[str]]:
340    """
341    Perform a split on a value and return N words as a result with `None` used for words that don't exist.
342
343    Args:
344        value: The value to be split.
345        sep: The value to use to split on.
346        min_num_words: The minimum number of words that are going to be in the result.
347        fill_from_start: Indicates that if `None` values should be inserted at the start or end of the list.
348
349    Examples:
350        >>> split_num_words("db.table", ".", 3)
351        [None, 'db', 'table']
352        >>> split_num_words("db.table", ".", 3, fill_from_start=False)
353        ['db', 'table', None]
354        >>> split_num_words("db.table", ".", 1)
355        ['db', 'table']
356
357    Returns:
358        The list of words returned by `split`, possibly augmented by a number of `None` values.
359    """
360    words = value.split(sep)
361    if fill_from_start:
362        return [None] * (min_num_words - len(words)) + words
363    return words + [None] * (min_num_words - len(words))
364
365
366def is_iterable(value: t.Any) -> bool:
367    """
368    Checks if the value is an iterable, excluding the types `str` and `bytes`.
369
370    Examples:
371        >>> is_iterable([1,2])
372        True
373        >>> is_iterable("test")
374        False
375
376    Args:
377        value: The value to check if it is an iterable.
378
379    Returns:
380        A `bool` value indicating if it is an iterable.
381    """
382    from sqlglot import Expression
383
384    return hasattr(value, "__iter__") and not isinstance(value, (str, bytes, Expression))
385
386
387def flatten(values: t.Iterable[t.Iterable[t.Any] | t.Any]) -> t.Iterator[t.Any]:
388    """
389    Flattens an iterable that can contain both iterable and non-iterable elements. Objects of
390    type `str` and `bytes` are not regarded as iterables.
391
392    Examples:
393        >>> list(flatten([[1, 2], 3, {4}, (5, "bla")]))
394        [1, 2, 3, 4, 5, 'bla']
395        >>> list(flatten([1, 2, 3]))
396        [1, 2, 3]
397
398    Args:
399        values: The value to be flattened.
400
401    Yields:
402        Non-iterable elements in `values`.
403    """
404    for value in values:
405        if is_iterable(value):
406            yield from flatten(value)
407        else:
408            yield value
409
410
411def dict_depth(d: t.Dict) -> int:
412    """
413    Get the nesting depth of a dictionary.
414
415    Example:
416        >>> dict_depth(None)
417        0
418        >>> dict_depth({})
419        1
420        >>> dict_depth({"a": "b"})
421        1
422        >>> dict_depth({"a": {}})
423        2
424        >>> dict_depth({"a": {"b": {}}})
425        3
426    """
427    try:
428        return 1 + dict_depth(next(iter(d.values())))
429    except AttributeError:
430        # d doesn't have attribute "values"
431        return 0
432    except StopIteration:
433        # d.values() returns an empty sequence
434        return 1
435
436
437def first(it: t.Iterable[T]) -> T:
438    """Returns the first element from an iterable (useful for sets)."""
439    return next(i for i in it)
440
441
442def merge_ranges(ranges: t.List[t.Tuple[A, A]]) -> t.List[t.Tuple[A, A]]:
443    """
444    Merges a sequence of ranges, represented as tuples (low, high) whose values
445    belong to some totally-ordered set.
446
447    Example:
448        >>> merge_ranges([(1, 3), (2, 6)])
449        [(1, 6)]
450    """
451    if not ranges:
452        return []
453
454    ranges = sorted(ranges)
455
456    merged = [ranges[0]]
457
458    for start, end in ranges[1:]:
459        last_start, last_end = merged[-1]
460
461        if start <= last_end:
462            merged[-1] = (last_start, max(last_end, end))
463        else:
464            merged.append((start, end))
465
466    return merged
467
468
469def is_iso_date(text: str) -> bool:
470    try:
471        datetime.date.fromisoformat(text)
472        return True
473    except ValueError:
474        return False
475
476
477def is_iso_datetime(text: str) -> bool:
478    try:
479        datetime.datetime.fromisoformat(text)
480        return True
481    except ValueError:
482        return False
483
484
485# Interval units that operate on date components
486DATE_UNITS = {"day", "week", "month", "quarter", "year", "year_month"}
487
488
489def is_date_unit(expression: t.Optional[exp.Expression]) -> bool:
490    return expression is not None and expression.name.lower() in DATE_UNITS