sqlglot.helper

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