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