Edit on GitHub

# sqlglot.executor.env

View Source

      1import datetime
      2import inspect
      3import re
      4import statistics
      5from functools import wraps
      6
      7from sqlglot import exp
      8from sqlglot.helper import PYTHON_VERSION
      9
     10
     11class reverse_key:
     12    def __init__(self, obj):
     13        self.obj = obj
     14
     15    def __eq__(self, other):
     16        return other.obj == self.obj
     17
     18    def __lt__(self, other):
     19        return other.obj < self.obj
     20
     21
     22def filter_nulls(func, empty_null=True):
     23    @wraps(func)
     24    def _func(values):
     25        filtered = tuple(v for v in values if v is not None)
     26        if not filtered and empty_null:
     27            return None
     28        return func(filtered)
     29
     30    return _func
     31
     32
     33def null_if_any(*required):
     34    """
     35    Decorator that makes a function return `None` if any of the `required` arguments are `None`.
     36
     37    This also supports decoration with no arguments, e.g.:
     38
     39        @null_if_any
     40        def foo(a, b): ...
     41
     42    In which case all arguments are required.
     43    """
     44    f = None
     45    if len(required) == 1 and callable(required[0]):
     46        f = required[0]
     47        required = ()
     48
     49    def decorator(func):
     50        if required:
     51            required_indices = [
     52                i for i, param in enumerate(inspect.signature(func).parameters) if param in required
     53            ]
     54
     55            def predicate(*args):
     56                return any(args[i] is None for i in required_indices)
     57
     58        else:
     59
     60            def predicate(*args):
     61                return any(a is None for a in args)
     62
     63        @wraps(func)
     64        def _func(*args):
     65            if predicate(*args):
     66                return None
     67            return func(*args)
     68
     69        return _func
     70
     71    if f:
     72        return decorator(f)
     73
     74    return decorator
     75
     76
     77@null_if_any("substr", "this")
     78def str_position(substr, this, position=None):
     79    position = position - 1 if position is not None else position
     80    return this.find(substr, position) + 1
     81
     82
     83@null_if_any("this")
     84def substring(this, start=None, length=None):
     85    if start is None:
     86        return this
     87    elif start == 0:
     88        return ""
     89    elif start < 0:
     90        start = len(this) + start
     91    else:
     92        start -= 1
     93
     94    end = None if length is None else start + length
     95
     96    return this[start:end]
     97
     98
     99@null_if_any
    100def cast(this, to):
    101    if to == exp.DataType.Type.DATE:
    102        return datetime.date.fromisoformat(this)
    103    if to == exp.DataType.Type.DATETIME:
    104        return datetime.datetime.fromisoformat(this)
    105    if to in exp.DataType.TEXT_TYPES:
    106        return str(this)
    107    if to in {exp.DataType.Type.FLOAT, exp.DataType.Type.DOUBLE}:
    108        return float(this)
    109    if to in exp.DataType.NUMERIC_TYPES:
    110        return int(this)
    111    raise NotImplementedError(f"Casting to '{to}' not implemented.")
    112
    113
    114def ordered(this, desc, nulls_first):
    115    if desc:
    116        return reverse_key(this)
    117    return this
    118
    119
    120@null_if_any
    121def interval(this, unit):
    122    if unit == "DAY":
    123        return datetime.timedelta(days=float(this))
    124    raise NotImplementedError
    125
    126
    127ENV = {
    128    "exp": exp,
    129    # aggs
    130    "ARRAYAGG": list,
    131    "AVG": filter_nulls(statistics.fmean if PYTHON_VERSION >= (3, 8) else statistics.mean),  # type: ignore
    132    "COUNT": filter_nulls(lambda acc: sum(1 for _ in acc), False),
    133    "MAX": filter_nulls(max),
    134    "MIN": filter_nulls(min),
    135    "SUM": filter_nulls(sum),
    136    # scalar functions
    137    "ABS": null_if_any(lambda this: abs(this)),
    138    "ADD": null_if_any(lambda e, this: e + this),
    139    "ARRAYANY": null_if_any(lambda arr, func: any(func(e) for e in arr)),
    140    "BETWEEN": null_if_any(lambda this, low, high: low <= this and this <= high),
    141    "BITWISEAND": null_if_any(lambda this, e: this & e),
    142    "BITWISELEFTSHIFT": null_if_any(lambda this, e: this << e),
    143    "BITWISEOR": null_if_any(lambda this, e: this | e),
    144    "BITWISERIGHTSHIFT": null_if_any(lambda this, e: this >> e),
    145    "BITWISEXOR": null_if_any(lambda this, e: this ^ e),
    146    "CAST": cast,
    147    "COALESCE": lambda *args: next((a for a in args if a is not None), None),
    148    "CONCAT": null_if_any(lambda *args: "".join(args)),
    149    "CONCATWS": null_if_any(lambda this, *args: this.join(args)),
    150    "DIV": null_if_any(lambda e, this: e / this),
    151    "EQ": null_if_any(lambda this, e: this == e),
    152    "EXTRACT": null_if_any(lambda this, e: getattr(e, this)),
    153    "GT": null_if_any(lambda this, e: this > e),
    154    "GTE": null_if_any(lambda this, e: this >= e),
    155    "IFNULL": lambda e, alt: alt if e is None else e,
    156    "IF": lambda predicate, true, false: true if predicate else false,
    157    "INTDIV": null_if_any(lambda e, this: e // this),
    158    "INTERVAL": interval,
    159    "LIKE": null_if_any(
    160        lambda this, e: bool(re.match(e.replace("_", ".").replace("%", ".*"), this))
    161    ),
    162    "LOWER": null_if_any(lambda arg: arg.lower()),
    163    "LT": null_if_any(lambda this, e: this < e),
    164    "LTE": null_if_any(lambda this, e: this <= e),
    165    "MOD": null_if_any(lambda e, this: e % this),
    166    "MUL": null_if_any(lambda e, this: e * this),
    167    "NEQ": null_if_any(lambda this, e: this != e),
    168    "ORD": null_if_any(ord),
    169    "ORDERED": ordered,
    170    "POW": pow,
    171    "STRPOSITION": str_position,
    172    "SUB": null_if_any(lambda e, this: e - this),
    173    "SUBSTRING": substring,
    174    "TIMESTRTOTIME": null_if_any(lambda arg: datetime.datetime.fromisoformat(arg)),
    175    "UPPER": null_if_any(lambda arg: arg.upper()),
    176}

class reverse_key: View Source

    12class reverse_key:
    13    def __init__(self, obj):
    14        self.obj = obj
    15
    16    def __eq__(self, other):
    17        return other.obj == self.obj
    18
    19    def __lt__(self, other):
    20        return other.obj < self.obj

reverse_key(obj) View Source

    13    def __init__(self, obj):
    14        self.obj = obj

def filter_nulls(func, empty_null=True): View Source

    23def filter_nulls(func, empty_null=True):
    24    @wraps(func)
    25    def _func(values):
    26        filtered = tuple(v for v in values if v is not None)
    27        if not filtered and empty_null:
    28            return None
    29        return func(filtered)
    30
    31    return _func

def null_if_any(*required): View Source

    34def null_if_any(*required):
    35    """
    36    Decorator that makes a function return `None` if any of the `required` arguments are `None`.
    37
    38    This also supports decoration with no arguments, e.g.:
    39
    40        @null_if_any
    41        def foo(a, b): ...
    42
    43    In which case all arguments are required.
    44    """
    45    f = None
    46    if len(required) == 1 and callable(required[0]):
    47        f = required[0]
    48        required = ()
    49
    50    def decorator(func):
    51        if required:
    52            required_indices = [
    53                i for i, param in enumerate(inspect.signature(func).parameters) if param in required
    54            ]
    55
    56            def predicate(*args):
    57                return any(args[i] is None for i in required_indices)
    58
    59        else:
    60
    61            def predicate(*args):
    62                return any(a is None for a in args)
    63
    64        @wraps(func)
    65        def _func(*args):
    66            if predicate(*args):
    67                return None
    68            return func(*args)
    69
    70        return _func
    71
    72    if f:
    73        return decorator(f)
    74
    75    return decorator

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.

@null_if_any('substr', 'this')
def str_position(substr, this, position=None): View Source

    78@null_if_any("substr", "this")
    79def str_position(substr, this, position=None):
    80    position = position - 1 if position is not None else position
    81    return this.find(substr, position) + 1

@null_if_any('this')
def substring(this, start=None, length=None): View Source

    84@null_if_any("this")
    85def substring(this, start=None, length=None):
    86    if start is None:
    87        return this
    88    elif start == 0:
    89        return ""
    90    elif start < 0:
    91        start = len(this) + start
    92    else:
    93        start -= 1
    94
    95    end = None if length is None else start + length
    96
    97    return this[start:end]

@null_if_any
def cast(this, to): View Source

    100@null_if_any
    101def cast(this, to):
    102    if to == exp.DataType.Type.DATE:
    103        return datetime.date.fromisoformat(this)
    104    if to == exp.DataType.Type.DATETIME:
    105        return datetime.datetime.fromisoformat(this)
    106    if to in exp.DataType.TEXT_TYPES:
    107        return str(this)
    108    if to in {exp.DataType.Type.FLOAT, exp.DataType.Type.DOUBLE}:
    109        return float(this)
    110    if to in exp.DataType.NUMERIC_TYPES:
    111        return int(this)
    112    raise NotImplementedError(f"Casting to '{to}' not implemented.")

def ordered(this, desc, nulls_first): View Source

    115def ordered(this, desc, nulls_first):
    116    if desc:
    117        return reverse_key(this)
    118    return this

@null_if_any
def interval(this, unit): View Source

    121@null_if_any
    122def interval(this, unit):
    123    if unit == "DAY":
    124        return datetime.timedelta(days=float(this))
    125    raise NotImplementedError