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