altimate-code 0.5.2 → 0.5.3

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/executor/context.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot.executor.env import ENV
+
+if t.TYPE_CHECKING:
+    from sqlglot.executor.table import Table, TableIter
+
+
+class Context:
+    """
+    Execution context for sql expressions.
+
+    Context is used to hold relevant data tables which can then be queried on with eval.
+
+    References to columns can either be scalar or vectors. When set_row is used, column references
+    evaluate to scalars while set_range evaluates to vectors. This allows convenient and efficient
+    evaluation of aggregation functions.
+    """
+
+    def __init__(self, tables: t.Dict[str, Table], env: t.Optional[t.Dict] = None) -> None:
+        """
+        Args
+            tables: representing the scope of the current execution context.
+            env: dictionary of functions within the execution context.
+        """
+        self.tables = tables
+        self._table: t.Optional[Table] = None
+        self.range_readers = {name: table.range_reader for name, table in self.tables.items()}
+        self.row_readers = {name: table.reader for name, table in tables.items()}
+        self.env = {**ENV, **(env or {}), "scope": self.row_readers}
+
+    def eval(self, code):
+        return eval(code, self.env)
+
+    def eval_tuple(self, codes):
+        return tuple(self.eval(code) for code in codes)
+
+    @property
+    def table(self) -> Table:
+        if self._table is None:
+            self._table = list(self.tables.values())[0]
+
+            for other in self.tables.values():
+                if self._table.columns != other.columns:
+                    raise Exception("Columns are different.")
+                if len(self._table.rows) != len(other.rows):
+                    raise Exception("Rows are different.")
+
+        return self._table
+
+    def add_columns(self, *columns: str) -> None:
+        for table in self.tables.values():
+            table.add_columns(*columns)
+
+    @property
+    def columns(self) -> t.Tuple:
+        return self.table.columns
+
+    def __iter__(self):
+        self.env["scope"] = self.row_readers
+        for i in range(len(self.table.rows)):
+            for table in self.tables.values():
+                reader = table[i]
+            yield reader, self
+
+    def table_iter(self, table: str) -> TableIter:
+        self.env["scope"] = self.row_readers
+        return iter(self.tables[table])
+
+    def filter(self, condition) -> None:
+        rows = [reader.row for reader, _ in self if self.eval(condition)]
+
+        for table in self.tables.values():
+            table.rows = rows
+
+    def sort(self, key) -> None:
+        def sort_key(row: t.Tuple) -> t.Tuple:
+            self.set_row(row)
+            return tuple((t is None, t) for t in self.eval_tuple(key))
+
+        self.table.rows.sort(key=sort_key)
+
+    def set_row(self, row: t.Tuple) -> None:
+        for table in self.tables.values():
+            table.reader.row = row
+        self.env["scope"] = self.row_readers
+
+    def set_index(self, index: int) -> None:
+        for table in self.tables.values():
+            table[index]
+        self.env["scope"] = self.row_readers
+
+    def set_range(self, start: int, end: int) -> None:
+        for name in self.tables:
+            self.range_readers[name].range = range(start, end)
+        self.env["scope"] = self.range_readers
+
+    def __contains__(self, table: str) -> bool:
+        return table in self.tables

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/executor/env.py

@@ -0,0 +1,246 @@
+import datetime
+import inspect
+import re
+import statistics
+from functools import wraps
+
+from sqlglot import exp
+from sqlglot.generator import Generator
+from sqlglot.helper import PYTHON_VERSION, is_int, seq_get
+
+
+class reverse_key:
+    def __init__(self, obj):
+        self.obj = obj
+
+    def __eq__(self, other):
+        return other.obj == self.obj
+
+    def __lt__(self, other):
+        return other.obj < self.obj
+
+
+def filter_nulls(func, empty_null=True):
+    @wraps(func)
+    def _func(values):
+        filtered = tuple(v for v in values if v is not None)
+        if not filtered and empty_null:
+            return None
+        return func(filtered)
+
+    return _func
+
+
+def null_if_any(*required):
+    """
+    Decorator that makes a function return `None` if any of the `required` arguments are `None`.
+
+    This also supports decoration with no arguments, e.g.:
+
+        @null_if_any
+        def foo(a, b): ...
+
+    In which case all arguments are required.
+    """
+    f = None
+    if len(required) == 1 and callable(required[0]):
+        f = required[0]
+        required = ()
+
+    def decorator(func):
+        if required:
+            required_indices = [
+                i for i, param in enumerate(inspect.signature(func).parameters) if param in required
+            ]
+
+            def predicate(*args):
+                return any(args[i] is None for i in required_indices)
+
+        else:
+
+            def predicate(*args):
+                return any(a is None for a in args)
+
+        @wraps(func)
+        def _func(*args):
+            if predicate(*args):
+                return None
+            return func(*args)
+
+        return _func
+
+    if f:
+        return decorator(f)
+
+    return decorator
+
+
+@null_if_any("this", "substr")
+def str_position(this, substr, position=None):
+    position = position - 1 if position is not None else position
+    return this.find(substr, position) + 1
+
+
+@null_if_any("this")
+def substring(this, start=None, length=None):
+    if start is None:
+        return this
+    elif start == 0:
+        return ""
+    elif start < 0:
+        start = len(this) + start
+    else:
+        start -= 1
+
+    end = None if length is None else start + length
+
+    return this[start:end]
+
+
+@null_if_any
+def cast(this, to):
+    if to == exp.DataType.Type.DATE:
+        if isinstance(this, datetime.datetime):
+            return this.date()
+        if isinstance(this, datetime.date):
+            return this
+        if isinstance(this, str):
+            return datetime.date.fromisoformat(this)
+    if to == exp.DataType.Type.TIME:
+        if isinstance(this, datetime.datetime):
+            return this.time()
+        if isinstance(this, datetime.time):
+            return this
+        if isinstance(this, str):
+            return datetime.time.fromisoformat(this)
+    if to in (exp.DataType.Type.DATETIME, exp.DataType.Type.TIMESTAMP):
+        if isinstance(this, datetime.datetime):
+            return this
+        if isinstance(this, datetime.date):
+            return datetime.datetime(this.year, this.month, this.day)
+        if isinstance(this, str):
+            return datetime.datetime.fromisoformat(this)
+    if to == exp.DataType.Type.BOOLEAN:
+        return bool(this)
+    if to in exp.DataType.TEXT_TYPES:
+        return str(this)
+    if to in {exp.DataType.Type.FLOAT, exp.DataType.Type.DOUBLE}:
+        return float(this)
+    if to in exp.DataType.NUMERIC_TYPES:
+        return int(this)
+    raise NotImplementedError(f"Casting {this} to '{to}' not implemented.")
+
+
+def ordered(this, desc, nulls_first):
+    if desc:
+        return reverse_key(this)
+    return this
+
+
+@null_if_any
+def interval(this, unit):
+    plural = unit + "S"
+    if plural in Generator.TIME_PART_SINGULARS:
+        unit = plural
+    return datetime.timedelta(**{unit.lower(): float(this)})
+
+
+@null_if_any("this", "expression")
+def arraytostring(this, expression, null=None):
+    return expression.join(x for x in (x if x is not None else null for x in this) if x is not None)
+
+
+@null_if_any("this", "expression")
+def jsonextract(this, expression):
+    for path_segment in expression:
+        if isinstance(this, dict):
+            this = this.get(path_segment)
+        elif isinstance(this, list) and is_int(path_segment):
+            this = seq_get(this, int(path_segment))
+        else:
+            raise NotImplementedError(f"Unable to extract value for {this} at {path_segment}.")
+
+        if this is None:
+            break
+
+    return this
+
+
+ENV = {
+    "exp": exp,
+    # aggs
+    "ARRAYAGG": list,
+    "ARRAYUNIQUEAGG": filter_nulls(lambda acc: list(set(acc))),
+    "AVG": filter_nulls(statistics.fmean if PYTHON_VERSION >= (3, 8) else statistics.mean),  # type: ignore
+    "COUNT": filter_nulls(lambda acc: sum(1 for _ in acc), False),
+    "MAX": filter_nulls(max),
+    "MIN": filter_nulls(min),
+    "SUM": filter_nulls(sum),
+    # scalar functions
+    "ABS": null_if_any(lambda this: abs(this)),
+    "ADD": null_if_any(lambda e, this: e + this),
+    "ARRAYANY": null_if_any(lambda arr, func: any(func(e) for e in arr)),
+    "ARRAYTOSTRING": arraytostring,
+    "BETWEEN": null_if_any(lambda this, low, high: low <= this and this <= high),
+    "BITWISEAND": null_if_any(lambda this, e: this & e),
+    "BITWISELEFTSHIFT": null_if_any(lambda this, e: this << e),
+    "BITWISEOR": null_if_any(lambda this, e: this | e),
+    "BITWISERIGHTSHIFT": null_if_any(lambda this, e: this >> e),
+    "BITWISEXOR": null_if_any(lambda this, e: this ^ e),
+    "CAST": cast,
+    "COALESCE": lambda *args: next((a for a in args if a is not None), None),
+    "CONCAT": null_if_any(lambda *args: "".join(args)),
+    "SAFECONCAT": null_if_any(lambda *args: "".join(str(arg) for arg in args)),
+    "CONCATWS": null_if_any(lambda this, *args: this.join(args)),
+    "DATEDIFF": null_if_any(lambda this, expression, *_: (this - expression).days),
+    "DATESTRTODATE": null_if_any(lambda arg: datetime.date.fromisoformat(arg)),
+    "DIV": null_if_any(lambda e, this: e / this),
+    "DOT": null_if_any(lambda e, this: e[this]),
+    "EQ": null_if_any(lambda this, e: this == e),
+    "EXTRACT": null_if_any(lambda this, e: getattr(e, this)),
+    "GT": null_if_any(lambda this, e: this > e),
+    "GTE": null_if_any(lambda this, e: this >= e),
+    "IF": lambda predicate, true, false: true if predicate else false,
+    "INTDIV": null_if_any(lambda e, this: e // this),
+    "INTERVAL": interval,
+    "JSONEXTRACT": jsonextract,
+    "LEFT": null_if_any(lambda this, e: this[:e]),
+    "LIKE": null_if_any(
+        lambda this, e: bool(re.match(e.replace("_", ".").replace("%", ".*"), this))
+    ),
+    "LOWER": null_if_any(lambda arg: arg.lower()),
+    "LT": null_if_any(lambda this, e: this < e),
+    "LTE": null_if_any(lambda this, e: this <= e),
+    "MAP": null_if_any(lambda *args: dict(zip(*args))),  # type: ignore
+    "MOD": null_if_any(lambda e, this: e % this),
+    "MUL": null_if_any(lambda e, this: e * this),
+    "NEQ": null_if_any(lambda this, e: this != e),
+    "ORD": null_if_any(ord),
+    "ORDERED": ordered,
+    "POW": pow,
+    "RIGHT": null_if_any(lambda this, e: this[-e:]),
+    "ROUND": null_if_any(lambda this, decimals=None, truncate=None: round(this, ndigits=decimals)),
+    "STRPOSITION": str_position,
+    "SUB": null_if_any(lambda e, this: e - this),
+    "SUBSTRING": substring,
+    "TIMESTRTOTIME": null_if_any(lambda arg: datetime.datetime.fromisoformat(arg)),
+    "UPPER": null_if_any(lambda arg: arg.upper()),
+    "YEAR": null_if_any(lambda arg: arg.year),
+    "MONTH": null_if_any(lambda arg: arg.month),
+    "DAY": null_if_any(lambda arg: arg.day),
+    "CURRENTDATETIME": datetime.datetime.now,
+    "CURRENTTIMESTAMP": datetime.datetime.now,
+    "CURRENTTIME": datetime.datetime.now,
+    "CURRENTDATE": datetime.date.today,
+    "STRFTIME": null_if_any(lambda fmt, arg: datetime.datetime.fromisoformat(arg).strftime(fmt)),
+    "STRTOTIME": null_if_any(lambda arg, format: datetime.datetime.strptime(arg, format)),
+    "TRIM": null_if_any(lambda this, e=None: this.strip(e)),
+    "STRUCT": lambda *args: {
+        args[x]: args[x + 1]
+        for x in range(0, len(args), 2)
+        if (args[x + 1] is not None and args[x] is not None)
+    },
+    "UNIXTOTIME": null_if_any(
+        lambda arg: datetime.datetime.fromtimestamp(arg, datetime.timezone.utc)
+    ),
+}