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