altimate-code 0.5.1 → 0.5.3

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/helper.py

@@ -0,0 +1,582 @@
+from __future__ import annotations
+
+import datetime
+import inspect
+import logging
+import re
+import sys
+import typing as t
+from collections.abc import Collection, Set
+from contextlib import contextmanager
+from copy import copy
+from difflib import get_close_matches
+from enum import Enum
+from itertools import count
+
+if t.TYPE_CHECKING:
+    from sqlglot import exp
+    from sqlglot._typing import A, E, T
+    from sqlglot.dialects.dialect import DialectType
+    from sqlglot.expressions import Expression
+
+
+CAMEL_CASE_PATTERN = re.compile("(?<!^)(?=[A-Z])")
+PYTHON_VERSION = sys.version_info[:2]
+logger = logging.getLogger("sqlglot")
+
+
+class AutoName(Enum):
+    """
+    This is used for creating Enum classes where `auto()` is the string form
+    of the corresponding enum's identifier (e.g. FOO.value results in "FOO").
+
+    Reference: https://docs.python.org/3/howto/enum.html#using-automatic-values
+    """
+
+    def _generate_next_value_(name, _start, _count, _last_values):
+        return name
+
+
+class classproperty(property):
+    """
+    Similar to a normal property but works for class methods
+    """
+
+    def __get__(self, obj: t.Any, owner: t.Any = None) -> t.Any:
+        return classmethod(self.fget).__get__(None, owner)()  # type: ignore
+
+
+def suggest_closest_match_and_fail(
+    kind: str,
+    word: str,
+    possibilities: t.Iterable[str],
+) -> None:
+    close_matches = get_close_matches(word, possibilities, n=1)
+
+    similar = seq_get(close_matches, 0) or ""
+    if similar:
+        similar = f" Did you mean {similar}?"
+
+    raise ValueError(f"Unknown {kind} '{word}'.{similar}")
+
+
+def seq_get(seq: t.Sequence[T], index: int) -> t.Optional[T]:
+    """Returns the value in `seq` at position `index`, or `None` if `index` is out of bounds."""
+    try:
+        return seq[index]
+    except IndexError:
+        return None
+
+
+@t.overload
+def ensure_list(value: t.Collection[T]) -> t.List[T]: ...
+
+
+@t.overload
+def ensure_list(value: None) -> t.List: ...
+
+
+@t.overload
+def ensure_list(value: T) -> t.List[T]: ...
+
+
+def ensure_list(value):
+    """
+    Ensures that a value is a list, otherwise casts or wraps it into one.
+
+    Args:
+        value: The value of interest.
+
+    Returns:
+        The value cast as a list if it's a list or a tuple, or else the value wrapped in a list.
+    """
+    if value is None:
+        return []
+    if isinstance(value, (list, tuple)):
+        return list(value)
+
+    return [value]
+
+
+@t.overload
+def ensure_collection(value: t.Collection[T]) -> t.Collection[T]: ...
+
+
+@t.overload
+def ensure_collection(value: T) -> t.Collection[T]: ...
+
+
+def ensure_collection(value):
+    """
+    Ensures that a value is a collection (excluding `str` and `bytes`), otherwise wraps it into a list.
+
+    Args:
+        value: The value of interest.
+
+    Returns:
+        The value if it's a collection, or else the value wrapped in a list.
+    """
+    if value is None:
+        return []
+    return (
+        value if isinstance(value, Collection) and not isinstance(value, (str, bytes)) else [value]
+    )
+
+
+def csv(*args: str, sep: str = ", ") -> str:
+    """
+    Formats any number of string arguments as CSV.
+
+    Args:
+        args: The string arguments to format.
+        sep: The argument separator.
+
+    Returns:
+        The arguments formatted as a CSV string.
+    """
+    return sep.join(arg for arg in args if arg)
+
+
+def subclasses(
+    module_name: str,
+    classes: t.Type | t.Tuple[t.Type, ...],
+    exclude: t.Type | t.Tuple[t.Type, ...] = (),
+) -> t.List[t.Type]:
+    """
+    Returns all subclasses for a collection of classes, possibly excluding some of them.
+
+    Args:
+        module_name: The name of the module to search for subclasses in.
+        classes: Class(es) we want to find the subclasses of.
+        exclude: Class(es) we want to exclude from the returned list.
+
+    Returns:
+        The target subclasses.
+    """
+    return [
+        obj
+        for _, obj in inspect.getmembers(
+            sys.modules[module_name],
+            lambda obj: inspect.isclass(obj) and issubclass(obj, classes) and obj not in exclude,
+        )
+    ]
+
+
+def apply_index_offset(
+    this: exp.Expression,
+    expressions: t.List[E],
+    offset: int,
+    dialect: DialectType = None,
+) -> t.List[E]:
+    """
+    Applies an offset to a given integer literal expression.
+
+    Args:
+        this: The target of the index.
+        expressions: The expression the offset will be applied to, wrapped in a list.
+        offset: The offset that will be applied.
+        dialect: the dialect of interest.
+
+    Returns:
+        The original expression with the offset applied to it, wrapped in a list. If the provided
+        `expressions` argument contains more than one expression, it's returned unaffected.
+    """
+    if not offset or len(expressions) != 1:
+        return expressions
+
+    expression = expressions[0]
+
+    from sqlglot import exp
+    from sqlglot.optimizer.annotate_types import annotate_types
+    from sqlglot.optimizer.simplify import simplify
+
+    if not this.type:
+        annotate_types(this, dialect=dialect)
+
+    if t.cast(exp.DataType, this.type).this not in (
+        exp.DataType.Type.UNKNOWN,
+        exp.DataType.Type.ARRAY,
+    ):
+        return expressions
+
+    if not expression.type:
+        annotate_types(expression, dialect=dialect)
+
+    if t.cast(exp.DataType, expression.type).this in exp.DataType.INTEGER_TYPES:
+        logger.info("Applying array index offset (%s)", offset)
+        expression = simplify(expression + offset)
+        return [expression]
+
+    return expressions
+
+
+def camel_to_snake_case(name: str) -> str:
+    """Converts `name` from camelCase to snake_case and returns the result."""
+    return CAMEL_CASE_PATTERN.sub("_", name).upper()
+
+
+def while_changing(expression: Expression, func: t.Callable[[Expression], E]) -> E:
+    """
+    Applies a transformation to a given expression until a fix point is reached.
+
+    Args:
+        expression: The expression to be transformed.
+        func: The transformation to be applied.
+
+    Returns:
+        The transformed expression.
+    """
+    end_hash: t.Optional[int] = None
+
+    while True:
+        # No need to walk the AST– we've already cached the hashes in the previous iteration
+        if end_hash is None:
+            for n in reversed(tuple(expression.walk())):
+                n._hash = hash(n)
+
+        start_hash = hash(expression)
+        expression = func(expression)
+
+        expression_nodes = tuple(expression.walk())
+
+        # Uncache previous caches so we can recompute them
+        for n in reversed(expression_nodes):
+            n._hash = None
+            n._hash = hash(n)
+
+        end_hash = hash(expression)
+
+        if start_hash == end_hash:
+            # ... and reset the hash so we don't risk it becoming out of date if a mutation happens
+            for n in expression_nodes:
+                n._hash = None
+
+            break
+
+    return expression
+
+
+def tsort(dag: t.Dict[T, t.Set[T]]) -> t.List[T]:
+    """
+    Sorts a given directed acyclic graph in topological order.
+
+    Args:
+        dag: The graph to be sorted.
+
+    Returns:
+        A list that contains all of the graph's nodes in topological order.
+    """
+    result = []
+
+    for node, deps in tuple(dag.items()):
+        for dep in deps:
+            if dep not in dag:
+                dag[dep] = set()
+
+    while dag:
+        current = {node for node, deps in dag.items() if not deps}
+
+        if not current:
+            raise ValueError("Cycle error")
+
+        for node in current:
+            dag.pop(node)
+
+        for deps in dag.values():
+            deps -= current
+
+        result.extend(sorted(current))  # type: ignore
+
+    return result
+
+
+def open_file(file_name: str) -> t.TextIO:
+    """Open a file that may be compressed as gzip and return it in universal newline mode."""
+    with open(file_name, "rb") as f:
+        gzipped = f.read(2) == b"\x1f\x8b"
+
+    if gzipped:
+        import gzip
+
+        return gzip.open(file_name, "rt", newline="")
+
+    return open(file_name, encoding="utf-8", newline="")
+
+
+@contextmanager
+def csv_reader(read_csv: exp.ReadCSV) -> t.Any:
+    """
+    Returns a csv reader given the expression `READ_CSV(name, ['delimiter', '|', ...])`.
+
+    Args:
+        read_csv: A `ReadCSV` function call.
+
+    Yields:
+        A python csv reader.
+    """
+    args = read_csv.expressions
+    file = open_file(read_csv.name)
+
+    delimiter = ","
+    args = iter(arg.name for arg in args)  # type: ignore
+    for k, v in zip(args, args):
+        if k == "delimiter":
+            delimiter = v
+
+    try:
+        import csv as csv_
+
+        yield csv_.reader(file, delimiter=delimiter)
+    finally:
+        file.close()
+
+
+def find_new_name(taken: t.Collection[str], base: str) -> str:
+    """
+    Searches for a new name.
+
+    Args:
+        taken: A collection of taken names.
+        base: Base name to alter.
+
+    Returns:
+        The new, available name.
+    """
+    if base not in taken:
+        return base
+
+    i = 2
+    new = f"{base}_{i}"
+    while new in taken:
+        i += 1
+        new = f"{base}_{i}"
+
+    return new
+
+
+def is_int(text: str) -> bool:
+    return is_type(text, int)
+
+
+def is_float(text: str) -> bool:
+    return is_type(text, float)
+
+
+def is_type(text: str, target_type: t.Type) -> bool:
+    try:
+        target_type(text)
+        return True
+    except ValueError:
+        return False
+
+
+def name_sequence(prefix: str) -> t.Callable[[], str]:
+    """Returns a name generator given a prefix (e.g. a0, a1, a2, ... if the prefix is "a")."""
+    sequence = count()
+    return lambda: f"{prefix}{next(sequence)}"
+
+
+def object_to_dict(obj: t.Any, **kwargs) -> t.Dict:
+    """Returns a dictionary created from an object's attributes."""
+    return {
+        **{k: v.copy() if hasattr(v, "copy") else copy(v) for k, v in vars(obj).items()},
+        **kwargs,
+    }
+
+
+def split_num_words(
+    value: str, sep: str, min_num_words: int, fill_from_start: bool = True
+) -> t.List[t.Optional[str]]:
+    """
+    Perform a split on a value and return N words as a result with `None` used for words that don't exist.
+
+    Args:
+        value: The value to be split.
+        sep: The value to use to split on.
+        min_num_words: The minimum number of words that are going to be in the result.
+        fill_from_start: Indicates that if `None` values should be inserted at the start or end of the list.
+
+    Examples:
+        >>> split_num_words("db.table", ".", 3)
+        [None, 'db', 'table']
+        >>> split_num_words("db.table", ".", 3, fill_from_start=False)
+        ['db', 'table', None]
+        >>> split_num_words("db.table", ".", 1)
+        ['db', 'table']
+
+    Returns:
+        The list of words returned by `split`, possibly augmented by a number of `None` values.
+    """
+    words = value.split(sep)
+    if fill_from_start:
+        return [None] * (min_num_words - len(words)) + words
+    return words + [None] * (min_num_words - len(words))
+
+
+def is_iterable(value: t.Any) -> bool:
+    """
+    Checks if the value is an iterable, excluding the types `str` and `bytes`.
+
+    Examples:
+        >>> is_iterable([1,2])
+        True
+        >>> is_iterable("test")
+        False
+
+    Args:
+        value: The value to check if it is an iterable.
+
+    Returns:
+        A `bool` value indicating if it is an iterable.
+    """
+    from sqlglot import Expression
+
+    return hasattr(value, "__iter__") and not isinstance(value, (str, bytes, Expression))
+
+
+def flatten(values: t.Iterable[t.Iterable[t.Any] | t.Any]) -> t.Iterator[t.Any]:
+    """
+    Flattens an iterable that can contain both iterable and non-iterable elements. Objects of
+    type `str` and `bytes` are not regarded as iterables.
+
+    Examples:
+        >>> list(flatten([[1, 2], 3, {4}, (5, "bla")]))
+        [1, 2, 3, 4, 5, 'bla']
+        >>> list(flatten([1, 2, 3]))
+        [1, 2, 3]
+
+    Args:
+        values: The value to be flattened.
+
+    Yields:
+        Non-iterable elements in `values`.
+    """
+    for value in values:
+        if is_iterable(value):
+            yield from flatten(value)
+        else:
+            yield value
+
+
+def dict_depth(d: t.Dict) -> int:
+    """
+    Get the nesting depth of a dictionary.
+
+    Example:
+        >>> dict_depth(None)
+        0
+        >>> dict_depth({})
+        1
+        >>> dict_depth({"a": "b"})
+        1
+        >>> dict_depth({"a": {}})
+        2
+        >>> dict_depth({"a": {"b": {}}})
+        3
+    """
+    try:
+        return 1 + dict_depth(next(iter(d.values())))
+    except AttributeError:
+        # d doesn't have attribute "values"
+        return 0
+    except StopIteration:
+        # d.values() returns an empty sequence
+        return 1
+
+
+def first(it: t.Iterable[T]) -> T:
+    """Returns the first element from an iterable (useful for sets)."""
+    return next(i for i in it)
+
+
+def to_bool(value: t.Optional[str | bool]) -> t.Optional[str | bool]:
+    if isinstance(value, bool) or value is None:
+        return value
+
+    # Coerce the value to boolean if it matches to the truthy/falsy values below
+    value_lower = value.lower()
+    if value_lower in ("true", "1"):
+        return True
+    if value_lower in ("false", "0"):
+        return False
+
+    return value
+
+
+def merge_ranges(ranges: t.List[t.Tuple[A, A]]) -> t.List[t.Tuple[A, A]]:
+    """
+    Merges a sequence of ranges, represented as tuples (low, high) whose values
+    belong to some totally-ordered set.
+
+    Example:
+        >>> merge_ranges([(1, 3), (2, 6)])
+        [(1, 6)]
+    """
+    if not ranges:
+        return []
+
+    ranges = sorted(ranges)
+
+    merged = [ranges[0]]
+
+    for start, end in ranges[1:]:
+        last_start, last_end = merged[-1]
+
+        if start <= last_end:
+            merged[-1] = (last_start, max(last_end, end))
+        else:
+            merged.append((start, end))
+
+    return merged
+
+
+def is_iso_date(text: str) -> bool:
+    try:
+        datetime.date.fromisoformat(text)
+        return True
+    except ValueError:
+        return False
+
+
+def is_iso_datetime(text: str) -> bool:
+    try:
+        datetime.datetime.fromisoformat(text)
+        return True
+    except ValueError:
+        return False
+
+
+# Interval units that operate on date components
+DATE_UNITS = {"day", "week", "month", "quarter", "year", "year_month"}
+
+
+def is_date_unit(expression: t.Optional[exp.Expression]) -> bool:
+    return expression is not None and expression.name.lower() in DATE_UNITS
+
+
+K = t.TypeVar("K")
+V = t.TypeVar("V")
+
+
+class SingleValuedMapping(t.Mapping[K, V]):
+    """
+    Mapping where all keys return the same value.
+
+    This rigamarole is meant to avoid copying keys, which was originally intended
+    as an optimization while qualifying columns for tables with lots of columns.
+    """
+
+    def __init__(self, keys: t.Collection[K], value: V):
+        self._keys = keys if isinstance(keys, Set) else set(keys)
+        self._value = value
+
+    def __getitem__(self, key: K) -> V:
+        if key in self._keys:
+            return self._value
+        raise KeyError(key)
+
+    def __len__(self) -> int:
+        return len(self._keys)
+
+    def __iter__(self) -> t.Iterator[K]:
+        return iter(self._keys)