sqlglot.executor.env
1import datetime 2import inspect 3import re 4import statistics 5from functools import wraps 6 7from sqlglot import exp 8from sqlglot.generator import Generator 9from sqlglot.helper import PYTHON_VERSION 10 11 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 21 22 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 32 33 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 76 77 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 82 83 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] 98 99 100@null_if_any 101def cast(this, to): 102 if to == exp.DataType.Type.DATE: 103 if isinstance(this, datetime.datetime): 104 return this.date() 105 if isinstance(this, datetime.date): 106 return this 107 if isinstance(this, str): 108 return datetime.date.fromisoformat(this) 109 if to in (exp.DataType.Type.DATETIME, exp.DataType.Type.TIMESTAMP): 110 if isinstance(this, datetime.datetime): 111 return this 112 if isinstance(this, datetime.date): 113 return datetime.datetime(this.year, this.month, this.day) 114 if isinstance(this, str): 115 return datetime.datetime.fromisoformat(this) 116 if to == exp.DataType.Type.BOOLEAN: 117 return bool(this) 118 if to in exp.DataType.TEXT_TYPES: 119 return str(this) 120 if to in {exp.DataType.Type.FLOAT, exp.DataType.Type.DOUBLE}: 121 return float(this) 122 if to in exp.DataType.NUMERIC_TYPES: 123 return int(this) 124 raise NotImplementedError(f"Casting {this} to '{to}' not implemented.") 125 126 127def ordered(this, desc, nulls_first): 128 if desc: 129 return reverse_key(this) 130 return this 131 132 133@null_if_any 134def interval(this, unit): 135 plural = unit + "S" 136 if plural in Generator.TIME_PART_SINGULARS: 137 unit = plural 138 return datetime.timedelta(**{unit.lower(): float(this)}) 139 140 141@null_if_any("this", "expression") 142def arrayjoin(this, expression, null=None): 143 return expression.join(x for x in (x if x is not None else null for x in this) if x is not None) 144 145 146ENV = { 147 "exp": exp, 148 # aggs 149 "ARRAYAGG": list, 150 "ARRAYUNIQUEAGG": filter_nulls(lambda acc: list(set(acc))), 151 "AVG": filter_nulls(statistics.fmean if PYTHON_VERSION >= (3, 8) else statistics.mean), # type: ignore 152 "COUNT": filter_nulls(lambda acc: sum(1 for _ in acc), False), 153 "MAX": filter_nulls(max), 154 "MIN": filter_nulls(min), 155 "SUM": filter_nulls(sum), 156 # scalar functions 157 "ABS": null_if_any(lambda this: abs(this)), 158 "ADD": null_if_any(lambda e, this: e + this), 159 "ARRAYANY": null_if_any(lambda arr, func: any(func(e) for e in arr)), 160 "ARRAYJOIN": arrayjoin, 161 "BETWEEN": null_if_any(lambda this, low, high: low <= this and this <= high), 162 "BITWISEAND": null_if_any(lambda this, e: this & e), 163 "BITWISELEFTSHIFT": null_if_any(lambda this, e: this << e), 164 "BITWISEOR": null_if_any(lambda this, e: this | e), 165 "BITWISERIGHTSHIFT": null_if_any(lambda this, e: this >> e), 166 "BITWISEXOR": null_if_any(lambda this, e: this ^ e), 167 "CAST": cast, 168 "COALESCE": lambda *args: next((a for a in args if a is not None), None), 169 "CONCAT": null_if_any(lambda *args: "".join(args)), 170 "SAFECONCAT": null_if_any(lambda *args: "".join(str(arg) for arg in args)), 171 "CONCATWS": null_if_any(lambda this, *args: this.join(args)), 172 "DATEDIFF": null_if_any(lambda this, expression, *_: (this - expression).days), 173 "DATESTRTODATE": null_if_any(lambda arg: datetime.date.fromisoformat(arg)), 174 "DIV": null_if_any(lambda e, this: e / this), 175 "DOT": null_if_any(lambda e, this: e[this]), 176 "EQ": null_if_any(lambda this, e: this == e), 177 "EXTRACT": null_if_any(lambda this, e: getattr(e, this)), 178 "GT": null_if_any(lambda this, e: this > e), 179 "GTE": null_if_any(lambda this, e: this >= e), 180 "IF": lambda predicate, true, false: true if predicate else false, 181 "INTDIV": null_if_any(lambda e, this: e // this), 182 "INTERVAL": interval, 183 "LEFT": null_if_any(lambda this, e: this[:e]), 184 "LIKE": null_if_any( 185 lambda this, e: bool(re.match(e.replace("_", ".").replace("%", ".*"), this)) 186 ), 187 "LOWER": null_if_any(lambda arg: arg.lower()), 188 "LT": null_if_any(lambda this, e: this < e), 189 "LTE": null_if_any(lambda this, e: this <= e), 190 "MAP": null_if_any(lambda *args: dict(zip(*args))), # type: ignore 191 "MOD": null_if_any(lambda e, this: e % this), 192 "MUL": null_if_any(lambda e, this: e * this), 193 "NEQ": null_if_any(lambda this, e: this != e), 194 "ORD": null_if_any(ord), 195 "ORDERED": ordered, 196 "POW": pow, 197 "RIGHT": null_if_any(lambda this, e: this[-e:]), 198 "STRPOSITION": str_position, 199 "SUB": null_if_any(lambda e, this: e - this), 200 "SUBSTRING": substring, 201 "TIMESTRTOTIME": null_if_any(lambda arg: datetime.datetime.fromisoformat(arg)), 202 "UPPER": null_if_any(lambda arg: arg.upper()), 203 "YEAR": null_if_any(lambda arg: arg.year), 204 "MONTH": null_if_any(lambda arg: arg.month), 205 "DAY": null_if_any(lambda arg: arg.day), 206 "CURRENTDATETIME": datetime.datetime.now, 207 "CURRENTTIMESTAMP": datetime.datetime.now, 208 "CURRENTTIME": datetime.datetime.now, 209 "CURRENTDATE": datetime.date.today, 210 "STRFTIME": null_if_any(lambda fmt, arg: datetime.datetime.fromisoformat(arg).strftime(fmt)), 211 "TRIM": null_if_any(lambda this, e=None: this.strip(e)), 212 "STRUCT": lambda *args: { 213 args[x]: args[x + 1] 214 for x in range(0, len(args), 2) 215 if (args[x + 1] is not None and args[x] is not None) 216 }, 217}
class
reverse_key:
def
filter_nulls(func, empty_null=True):
def
null_if_any(*required):
35def null_if_any(*required): 36 """ 37 Decorator that makes a function return `None` if any of the `required` arguments are `None`. 38 39 This also supports decoration with no arguments, e.g.: 40 41 @null_if_any 42 def foo(a, b): ... 43 44 In which case all arguments are required. 45 """ 46 f = None 47 if len(required) == 1 and callable(required[0]): 48 f = required[0] 49 required = () 50 51 def decorator(func): 52 if required: 53 required_indices = [ 54 i for i, param in enumerate(inspect.signature(func).parameters) if param in required 55 ] 56 57 def predicate(*args): 58 return any(args[i] is None for i in required_indices) 59 60 else: 61 62 def predicate(*args): 63 return any(a is None for a in args) 64 65 @wraps(func) 66 def _func(*args): 67 if predicate(*args): 68 return None 69 return func(*args) 70 71 return _func 72 73 if f: 74 return decorator(f) 75 76 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):
@null_if_any('this')
def
substring(this, start=None, length=None):
@null_if_any
def
cast(this, to):
101@null_if_any 102def cast(this, to): 103 if to == exp.DataType.Type.DATE: 104 if isinstance(this, datetime.datetime): 105 return this.date() 106 if isinstance(this, datetime.date): 107 return this 108 if isinstance(this, str): 109 return datetime.date.fromisoformat(this) 110 if to in (exp.DataType.Type.DATETIME, exp.DataType.Type.TIMESTAMP): 111 if isinstance(this, datetime.datetime): 112 return this 113 if isinstance(this, datetime.date): 114 return datetime.datetime(this.year, this.month, this.day) 115 if isinstance(this, str): 116 return datetime.datetime.fromisoformat(this) 117 if to == exp.DataType.Type.BOOLEAN: 118 return bool(this) 119 if to in exp.DataType.TEXT_TYPES: 120 return str(this) 121 if to in {exp.DataType.Type.FLOAT, exp.DataType.Type.DOUBLE}: 122 return float(this) 123 if to in exp.DataType.NUMERIC_TYPES: 124 return int(this) 125 raise NotImplementedError(f"Casting {this} to '{to}' not implemented.")
def
ordered(this, desc, nulls_first):
@null_if_any
def
interval(this, unit):
@null_if_any('this', 'expression')
def
arrayjoin(this, expression, null=None):
ENV =
{'exp': <module 'sqlglot.expressions' from '/home/runner/work/sqlglot/sqlglot/sqlglot/expressions.py'>, 'ARRAYAGG': <class 'list'>, 'ARRAYUNIQUEAGG': <function <lambda>>, 'AVG': <function fmean>, 'COUNT': <function <lambda>>, 'MAX': <function max>, 'MIN': <function min>, 'SUM': <function sum>, 'ABS': <function <lambda>>, 'ADD': <function <lambda>>, 'ARRAYANY': <function <lambda>>, 'ARRAYJOIN': <function arrayjoin>, 'BETWEEN': <function <lambda>>, 'BITWISEAND': <function <lambda>>, 'BITWISELEFTSHIFT': <function <lambda>>, 'BITWISEOR': <function <lambda>>, 'BITWISERIGHTSHIFT': <function <lambda>>, 'BITWISEXOR': <function <lambda>>, 'CAST': <function cast>, 'COALESCE': <function <lambda>>, 'CONCAT': <function <lambda>>, 'SAFECONCAT': <function <lambda>>, 'CONCATWS': <function <lambda>>, 'DATEDIFF': <function <lambda>>, 'DATESTRTODATE': <function <lambda>>, 'DIV': <function <lambda>>, 'DOT': <function <lambda>>, 'EQ': <function <lambda>>, 'EXTRACT': <function <lambda>>, 'GT': <function <lambda>>, 'GTE': <function <lambda>>, 'IF': <function <lambda>>, 'INTDIV': <function <lambda>>, 'INTERVAL': <function interval>, 'LEFT': <function <lambda>>, 'LIKE': <function <lambda>>, 'LOWER': <function <lambda>>, 'LT': <function <lambda>>, 'LTE': <function <lambda>>, 'MAP': <function <lambda>>, 'MOD': <function <lambda>>, 'MUL': <function <lambda>>, 'NEQ': <function <lambda>>, 'ORD': <function ord>, 'ORDERED': <function ordered>, 'POW': <built-in function pow>, 'RIGHT': <function <lambda>>, 'STRPOSITION': <function str_position>, 'SUB': <function <lambda>>, 'SUBSTRING': <function substring>, 'TIMESTRTOTIME': <function <lambda>>, 'UPPER': <function <lambda>>, 'YEAR': <function <lambda>>, 'MONTH': <function <lambda>>, 'DAY': <function <lambda>>, 'CURRENTDATETIME': <built-in method now of type object>, 'CURRENTTIMESTAMP': <built-in method now of type object>, 'CURRENTTIME': <built-in method now of type object>, 'CURRENTDATE': <built-in method today of type object>, 'STRFTIME': <function <lambda>>, 'TRIM': <function <lambda>>, 'STRUCT': <function <lambda>>}