ygo 1.0.3__py3-none-any.whl → 1.0.4__py3-none-any.whl

qdf/expr.py

@@ -0,0 +1,301 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Created on 2025/3/3 19:52
+@author: ZhangYundi
+@email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+
+from __future__ import annotations
+
+import re
+import warnings
+from dataclasses import dataclass
+
+from lark import Lark, Transformer, v_args
+
+from .errors import ParseError
+
+
+# 基类
+class Token:
+    pass
+
+
+@dataclass
+class OperatorToken(Token):
+    """算子类型token"""
+    value: str
+
+
+@dataclass
+class OperandToken(Token):
+    """运算对象token"""
+    value: str | float | int
+
+
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore")
+    grammar = """
+        start: expr
+        ?expr: ternary_expr
+        ?ternary_expr: or_expr
+             | or_expr "?" or_expr ":" ternary_expr -> ternary
+        ?or_expr: and_expr
+             | or_expr "|" and_expr -> or_
+        ?and_expr: comp_expr
+             | and_expr "&" comp_expr -> and_
+        ?comp_expr: eq_expr
+             | comp_expr "<" eq_expr -> lt
+             | comp_expr ">" eq_expr -> gt
+             | comp_expr "<=" eq_expr -> le
+             | comp_expr ">=" eq_expr -> ge
+        ?eq_expr: arith_expr
+             | eq_expr "==" arith_expr -> eq
+             | eq_expr "!=" arith_expr -> neq
+        ?arith_expr: term
+             | arith_expr "+" term -> add
+             | arith_expr "-" term -> sub
+        ?term: pow_expr
+             | term "*" pow_expr -> mul
+             | term "/" pow_expr -> div
+             | term "//" pow_expr -> floordiv // 取整
+             | term "%" pow_expr -> mod      // 求余
+        ?pow_expr: factor
+             | factor "**" pow_expr -> pow
+        ?factor: atom
+             | "-" factor -> neg
+             | "!" factor -> not_
+             | "~" factor -> not_
+        ?atom: function
+             | NAME
+             | NUMBER
+             | FLOAT
+             | "(" expr ")"
+             | implicit_mul // 隐式乘法
+             | attribute_access // 新增：属性访问
+        implicit_mul: (NUMBER | FLOAT) NAME -> implicit_mul  // 隐式乘法
+        attribute_access: atom "." NAME -> attribute_access // 新增：属性访问
+        function: NAME "(" expr_list ")" -> function
+        // expr_list: expr ("," expr)*
+        keyword_arg: NAME "=" expr -> keyword_arg  // 关键字参数
+        expr_list: (expr | keyword_arg) ("," (expr | keyword_arg))*  // 支持关键字参数
+        NAME: /[a-zA-Z_$,][a-zA-Z0-9_$]*/
+        NUMBER: /\d+/  // regex for numbers
+        FLOAT: /\d+\.\d+/
+        %import common.WS
+        %ignore WS
+    """
+
+
+class ExprParser(Transformer):
+    @v_args(inline=True)
+    def ternary(self, a, b, c):
+        return Expr.new("if_", [a, b, c])
+
+    def attribute_access(self, items):
+        return ".".join(items)
+
+    def keyword_arg(self, item):
+        k, v = item
+        return {k: v}
+
+    def NAME(self, name):
+        return str(name)
+
+    def NUMBER(self, number):  # new transformer for numbers
+        return int(number)
+
+    def FLOAT(self, number):
+        return float(number)
+
+    def add(self, items):
+        return Expr.new("add", items)
+
+    def sub(self, items):
+        return Expr.new("sub", items)
+
+    def mul(self, items):
+        return Expr.new("mul", items)
+
+    def div(self, items):
+        return Expr.new("div", items)
+
+    def floordiv(self, items):
+        return Expr.new("floordiv", items)
+
+    def mod(self, items):
+        return Expr.new("mod", items)
+
+    def pow(self, items):
+        return Expr.new("pow", items)
+
+    def neg(self, items):
+        item = items[0]
+        if isinstance(item, (int, float)):
+            return -item
+        return Expr.new("neg", items)
+
+    def not_(self, item):
+        return Expr.new("not_", item)
+
+    def and_(self, items):
+        return Expr.new("and_", items)
+
+    def or_(self, items):
+        return Expr.new("or_", items)
+
+    def eq(self, items):
+        return Expr.new("eq", items)
+
+    def neq(self, items):
+        return Expr.new("neq", items)
+
+    def lt(self, items):
+        return Expr.new("lt", items)
+
+    def gt(self, items):
+        return Expr.new("gt", items)
+
+    def le(self, items):
+        return Expr.new("le", items)
+
+    def ge(self, items):
+        return Expr.new("ge", items)
+
+    def function(self, items):
+        name = items.pop(0)
+        return Expr.new(name, items[0])
+
+    def implicit_mul(self, items):
+        return Expr.new("mul", items)
+
+    def expr_list(self, items):
+        return items
+
+
+parser = Lark(grammar, parser='lalr', transformer=ExprParser())
+
+
+def parse_expr(expression: str) -> Expr:
+    return parser.parse(expression).children[0]
+
+
+class Expr:
+
+    def __init__(self, expr: str | None = None):
+
+        self.fn_name: str | None = ""
+        self.args: list | None = None
+        self.alias: str | None = None
+        if expr:
+            try:
+                self._parse(expr)
+            except Exception as e:
+                raise ParseError(f"{expr}\n{e}")
+
+    @classmethod
+    def new(cls, fn_name: str | None, args: list | None, alias: str | None = None):
+        expr = cls()
+        expr.fn_name = fn_name
+        expr.args = args
+        expr.alias = alias if alias is not None else str(expr)
+        return expr
+
+    def to_rpn(self) -> list[Token]:
+        """生成逆波兰表达式: (后缀表达式: 运算符在后)"""
+        rpn = list()
+
+        # 递归遍历子表达式
+        def _traverse(node: Expr):
+
+            if node.args is not None:
+                for child in node.args:
+                    if isinstance(child, Expr):
+                        _traverse(child)
+                    else:
+                        rpn.append(OperandToken(child))
+            rpn.append(OperatorToken(node.fn_name))
+
+        _traverse(self)
+
+        return rpn
+
+    def __str__(self):
+        unary_map = {"neg": "-", "not_": "!"}
+        binary_map = {"add": "+",
+                      "mul": "*",
+                      "div": "/",
+                      "sub": "-",
+                      "floordiv": "//",
+                      "mod": "%",
+                      "pow": "**",
+                      "and_": "&",
+                      "or_": "|",
+                      "gt": ">",
+                      "gte": ">=",
+                      "lt": "<",
+                      "lte": "<=",
+                      "eq": "==",
+                      "neq": "!=",
+                      }
+        if self.fn_name is None:
+            return str(self.args[0])
+        if self.fn_name == "if_":
+            cond, body, orelse = self.args
+            return f"{cond}?{body}:{orelse}"
+        elif self.fn_name in ("neg", "not_"):
+            return f"{unary_map.get(self.fn_name)}{self.args[0]}"
+        elif self.fn_name in binary_map:
+            return f"({binary_map.get(self.fn_name).join([str(arg) for arg in self.args])})"
+        else:
+            return f"{self.fn_name}({', '.join([str(arg) for arg in self.args])})"
+
+    def __repr__(self):
+        return self.__str__()
+
+    def _parse(self, expr):
+        """
+        解析表达式
+        """
+        convertor = {
+            'if(': 'if_(',
+            'not(': 'not_(',
+            'and(': 'and_(',
+            'or(': 'or_(',
+            '$': '',
+            "\n": '',
+            "!": "~",
+            ",": ", ",
+        }
+        for old, new in convertor.items():
+            expr = expr.replace(old, new)
+        new_expr = expr
+        match = re.search(r'(?i)(.+?)\s+AS\s+(\w+)', new_expr)
+        alias = None
+        if match:
+            new_expr = match.group(1).strip()
+            alias = match.group(2).strip()
+
+        expr_ = parse_expr(new_expr)
+        self.alias = alias if alias is not None else str(expr_)
+        if not isinstance(expr_, Expr):
+            self.args = [expr_]
+        else:
+            self.fn_name, self.args = expr_.fn_name, expr_.args
+
+    @property
+    def n_args(self) -> int:
+        """返回表达式的参数个数"""
+        return len(self.args)
+
+    @property
+    def depth(self) -> int:
+        """返回表达式的嵌套深度"""
+        _depth = 1
+        _depths = [0]
+        for arg in self.args:
+            if isinstance(arg, Expr):
+                _depths.append(arg.depth)
+        return _depth + max(_depths)

qdf/qdf.py

@@ -0,0 +1,158 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Created on 2025/3/5 21:40
+@author: ZhangYundi
+@email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+from __future__ import annotations
+
+import importlib.util
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+import polars as pl
+from toolz import partial
+from functools import lru_cache
+
+import ygo
+import ylog
+from .errors import CalculateError, PolarsError
+from .expr import Expr
+import time
+
+# 动态加载模块
+module_name = "udf"
+module_path = Path(__file__).parent / "udf" / "__init__.py"
+spec = importlib.util.spec_from_file_location(module_name, module_path)
+module = importlib.util.module_from_spec(spec)
+sys.modules[module_name] = module
+spec.loader.exec_module(module)
+
+
+@dataclass
+class FailInfo:
+    expr: str
+    error: Exception
+
+    def __str__(self):
+        return f"""
+expr={self.expr}
+=================================================
+{self.error}
+=================================================
+"""
+
+    def __repr__(self):
+        return self.__str__()
+
+@lru_cache(maxsize=512)
+def parse_expr(expr: str) -> Expr:
+    return Expr(expr)
+
+class QDF:
+
+    def __init__(self,
+                 data: pl.LazyFrame,
+                 index: tuple[str] = ("date", "time", "asset"),
+                 align: bool = True,):
+        self.data = data.with_columns(pl.col(pl.Decimal).cast(pl.Float32))
+        self.dims = [self.data.select(index_).drop_nulls().unique().count().collect().item() for index_ in index]
+        if align:
+            lev_vals: list[pl.DataFrame] = [self.data.select(name).drop_nulls().unique() for name in index]
+            full_index = lev_vals[0]
+            for lev_val in lev_vals[1:]:
+                full_index = full_index.join(lev_val, how="cross")
+            self.data = full_index.join(self.data, on=index, how='left') #.sort(index).collect().lazy()
+
+        self.index = index
+        self.failed = list()
+
+    def __str__(self):
+        return self.data.__str__()
+
+    def __repr__(self):
+        return self.data.__str__()
+
+    def register_udf(self, func: callable, name: str = None):
+        name = name if name is not None else func.__name__
+        setattr(module, name, func)
+
+    def _compile_expr(self, expr: str, cover: bool):
+        expr_parsed = Expr(expr)
+        alias = expr_parsed.alias  # if expr_parsed.alias is not None else str(expr_parsed)
+        current_cols = set(self.data.collect_schema().keys())
+        columns = self.data.collect_schema().names()
+        if alias in current_cols and not cover:
+            return alias
+
+        def calc(expr_: Expr):
+            alias_ = expr_.alias
+            # _cols = self.data.collect_schema().names()
+            if alias_ in current_cols and not cover:
+                # 已存在：直接select数据源
+                return alias_
+            func = getattr(module, expr_.fn_name)
+            _params = ygo.fn_signature_params(func)
+            if "dims" in _params:
+                func = partial(func, dims=self.dims)
+            args = list()
+            kwargs = dict()
+            for arg in expr_.args:
+                if isinstance(arg, Expr):
+                    args.append(pl.col(calc(arg)))
+                elif isinstance(arg, dict):
+                    kwargs.update(arg)
+                elif isinstance(arg, str):
+                    args.append(pl.col(arg))
+                else:
+                    args.append(arg)  # or args.append(pl.lit(arg))
+            try:
+                expr_pl: pl.Expr = func(*args, **kwargs).alias(alias_)
+            except Exception as e:
+                raise CalculateError(f"{expr_.fn_name}({', '.join([str(arg) for arg in args])})\n{e}")
+            try:
+                self.data = self.data.with_columns(expr_pl)
+            except Exception as e:
+                raise PolarsError(f"{expr_}\n{e}")
+            return alias_
+
+        calc(expr_parsed)
+
+        columns.append(alias)
+        drop = current_cols.difference(set(columns))
+        self.data = self.data.drop(*drop)
+
+        return alias
+
+    def sql(self, *exprs: str, cover: bool = False,) -> pl.LazyFrame:
+        """
+        表达式查询
+        Parameters
+        ----------
+        exprs: str
+            表达式，比如 "ts_mean(close, 5) as close_ma5"
+        cover: bool
+            当遇到已经存在列名的时候，是否重新计算覆盖原来的列, 默认False，返回已经存在的列，跳过计算
+            - True: 重新计算并且返回新的结果，覆盖掉原来的列
+            - False, 返回已经存在的列，跳过计算
+        Returns
+        -------
+            polars.DataFrame
+        """
+        self.failed = list()
+        exprs_to_add = list()
+        for expr in exprs:
+            try:
+                compiled = self._compile_expr(expr, cover)
+                if compiled is not None:
+                    exprs_to_add.append(compiled)
+            except Exception as e:
+                self.failed.append(FailInfo(expr, e))
+        if self.failed:
+            ylog.warning(f"QDF.sql 失败：{len(self.failed)}/{len(exprs)}: \n {self.failed}")
+        final_df = self.data.with_columns(exprs_to_add).select(*self.index, *exprs_to_add).fill_nan(None).drop_nulls().sort(self.index)
+        return final_df.collect()
+

qdf/udf/__init__.py

@@ -0,0 +1,14 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Created on 2025/3/4 20:20
+@author: ZhangYundi
+@email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+
+from .base_udf import *
+from .cs_udf import *
+from .ts_udf import *
+from .d_udf import *
+from .ind_udf import *

qdf/udf/base_udf.py

@@ -0,0 +1,145 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Created on 2025/3/4 20:28
+@author: ZhangYundi
+@email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+
+import polars as pl
+import math
+
+"""
+基本算子：一元算子、二元算子、三元算子 以及 polars 支持的表达式(剔除数据泄露的)
+"""
+# ======================== 一元算子 ========================
+
+def not_(expr: pl.Expr): return ~expr
+
+
+def neg(expr: pl.Expr): return -expr
+
+
+def abs(expr: pl.Expr): return expr.abs()
+
+
+def log(expr: pl.Expr, base=math.e): return expr.log(base=base)
+
+
+def sqrt(expr: pl.Expr): return expr.sqrt()
+
+
+def square(expr: pl.Expr): return expr ** 2
+
+
+def cube(expr: pl.Expr): return expr ** 3
+
+
+def cbrt(expr: pl.Expr): return expr ** (1 / 3)
+
+
+def sin(expr: pl.Expr): return expr.sin()
+
+def sinh(expr: pl.Expr): return expr.sinh()
+
+def arcsin(expr: pl.Expr): return expr.arcsin()
+
+def arcsinh(expr: pl.Expr): return expr.arcsinh()
+
+
+def cos(expr: pl.Expr): return expr.cos()
+
+def cosh(expr: pl.Expr): return expr.cosh()
+
+def arccos(expr: pl.Expr): return expr.arccos()
+
+def arccosh(expr: pl.Expr): return expr.arccosh()
+
+def tan(expr: pl.Expr): return expr.tan()
+
+def tanh(expr: pl.Expr): return expr.tanh()
+
+def arctan(expr: pl.Expr): return expr.arctan()
+
+def arctanh(expr: pl.Expr): return expr.arctanh()
+
+
+def sign(expr: pl.Expr): return expr.sign()
+
+
+def sigmoid(expr: pl.Expr): return 1 / (1 + (-expr).exp())
+
+
+# def all(expr: pl.Expr, ignore_nulls: bool = True): return expr.all(ignore_nulls=ignore_nulls)
+
+
+# def any(expr: pl.Expr, ignore_nulls: bool = True): return expr.any(ignore_nulls=ignore_nulls)
+
+def cot(expr: pl.Expr): return expr.cot()
+
+def degrees(expr: pl.Expr): return expr.degrees()
+
+def exp(expr: pl.Expr): return expr.exp()
+
+def log1p(expr: pl.Expr): return expr.log1p()
+
+def clip(expr: pl.Expr, lower_bound, upper_bound): return expr.clip(lower_bound, upper_bound)
+
+# ======================== 二元算子 ========================
+def add(left: pl.Expr, right: pl.Expr): return left + right
+
+
+def sub(left: pl.Expr, right: pl.Expr): return left - right
+
+
+def mul(left: pl.Expr, right: pl.Expr): return left * right
+
+
+def div(left: pl.Expr, right: pl.Expr): return left / right
+
+
+def floordiv(left: pl.Expr, right: pl.Expr): return left // right
+
+
+def mod(left: pl.Expr, right: pl.Expr): return left % right
+
+
+def lt(left: pl.Expr, right: pl.Expr): return left < right
+
+
+def le(left: pl.Expr, right: pl.Expr): return left <= right
+
+
+def gt(left: pl.Expr, right: pl.Expr): return left > right
+
+
+def ge(left: pl.Expr, right: pl.Expr): return left >= right
+
+
+def eq(left: pl.Expr, right: pl.Expr): return left == right
+
+
+def neq(left: pl.Expr, right: pl.Expr): return left != right
+
+def and_(left: pl.Expr, right: pl.Expr): return left & right
+
+def or_(left: pl.Expr, right: pl.Expr): return left | right
+
+def max(*exprs: pl.Expr): return pl.max_horizontal(*exprs)
+
+def min(*exprs: pl.Expr): return pl.min_horizontal(*exprs)
+
+def sum(*exprs: pl.Expr): return pl.sum_horizontal(*exprs)
+
+
+# ======================== 三元 ========================
+def if_(cond: pl.Expr, body: pl.Expr, or_else: pl.Expr):
+    return pl.when(cond).then(body).otherwise(or_else)
+
+def fib(high: pl.Expr, low: pl.Expr, ratio: float = 0.618):
+    """
+    计算裴波那契回调比率
+    ratio: 0.236 | 0.382 | 0.618 等黄金分割比例
+    """
+    return low + (high - low) * ratio