quantnodes 3.0.0__py3-none-any.whl

QuantNodes/symbolic/expression.py

@@ -0,0 +1,234 @@
+# coding=utf-8
+"""
+符号计算引擎 - SQL 表达式节点
+
+定义用于 SQL 编译的表达式 AST 节点。
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, List, Optional, Tuple
+
+
+class SQLExpression(ABC):
+    """SQL 表达式基类"""
+
+    @abstractmethod
+    def to_sql(self, dialect: "SQLDialect") -> str:  # noqa: F821
+        """编译为 SQL 字符串"""
+        pass
+
+    def __add__(self, other: Any) -> "SQLBinaryOp":
+        return SQLBinaryOp(self, "+", wrap_sql_expr(other))
+
+    def __radd__(self, other: Any) -> "SQLBinaryOp":
+        return SQLBinaryOp(wrap_sql_expr(other), "+", self)
+
+    def __sub__(self, other: Any) -> "SQLBinaryOp":
+        return SQLBinaryOp(self, "-", wrap_sql_expr(other))
+
+    def __rsub__(self, other: Any) -> "SQLBinaryOp":
+        return SQLBinaryOp(wrap_sql_expr(other), "-", self)
+
+    def __mul__(self, other: Any) -> "SQLBinaryOp":
+        return SQLBinaryOp(self, "*", wrap_sql_expr(other))
+
+    def __rmul__(self, other: Any) -> "SQLBinaryOp":
+        return SQLBinaryOp(wrap_sql_expr(other), "*", self)
+
+    def __truediv__(self, other: Any) -> "SQLBinaryOp":
+        return SQLBinaryOp(self, "/", wrap_sql_expr(other))
+
+    def __rtruediv__(self, other: Any) -> "SQLBinaryOp":
+        return SQLBinaryOp(wrap_sql_expr(other), "/", self)
+
+    def __neg__(self) -> "SQLUnaryOp":
+        return SQLUnaryOp("-", self)
+
+    def __gt__(self, other: Any) -> "SQLComparison":
+        return SQLComparison(self, ">", wrap_sql_expr(other))
+
+    def __ge__(self, other: Any) -> "SQLComparison":
+        return SQLComparison(self, ">=", wrap_sql_expr(other))
+
+    def __lt__(self, other: Any) -> "SQLComparison":
+        return SQLComparison(self, "<", wrap_sql_expr(other))
+
+    def __le__(self, other: Any) -> "SQLComparison":
+        return SQLComparison(self, "<=", wrap_sql_expr(other))
+
+    def __eq__(self, other: Any) -> "SQLComparison":
+        return SQLComparison(self, "==", wrap_sql_expr(other))
+
+    def __ne__(self, other: Any) -> "SQLComparison":
+        return SQLComparison(self, "!=", wrap_sql_expr(other))
+
+    def __and__(self, other: Any) -> "SQLLogicalOp":
+        return SQLLogicalOp("AND", self, wrap_sql_expr(other))
+
+    def __or__(self, other: Any) -> "SQLLogicalOp":
+        return SQLLogicalOp("OR", self, wrap_sql_expr(other))
+
+    def __invert__(self) -> "SQLLogicalOp":
+        return SQLLogicalOp("NOT", self)
+
+
+class ColumnRef(SQLExpression):
+    """列引用表达式"""
+
+    def __init__(self, name: str, table: Optional[str] = None):
+        self.name = name
+        self.table = table
+
+    def to_sql(self, dialect) -> str:
+        if self.table:
+            return f"{dialect.quote_identifier(self.table)}.{dialect.quote_identifier(self.name)}"
+        return dialect.quote_identifier(self.name)
+
+    def __repr__(self) -> str:
+        if self.table:
+            return f"{self.table}.{self.name}"
+        return self.name
+
+
+class LiteralValue(SQLExpression):
+    """字面量值表达式"""
+
+    def __init__(self, value: Any):
+        self.value = value
+
+    def to_sql(self, dialect) -> str:
+        return dialect.quote_literal(self.value)
+
+    def __repr__(self) -> str:
+        return repr(self.value)
+
+
+class SQLBinaryOp(SQLExpression):
+    """二元运算表达式"""
+
+    def __init__(self, left: SQLExpression, op: str, right: SQLExpression):
+        self.left = left
+        self.op = op
+        self.right = right
+
+    def to_sql(self, dialect) -> str:
+        left_sql = self.left.to_sql(dialect)
+        right_sql = self.right.to_sql(dialect)
+        return f"({left_sql} {self.op} {right_sql})"
+
+    def __repr__(self) -> str:
+        return f"({self.left} {self.op} {self.right})"
+
+
+class SQLUnaryOp(SQLExpression):
+    """一元运算表达式"""
+
+    def __init__(self, op: str, operand: SQLExpression):
+        self.op = op
+        self.operand = operand
+
+    def to_sql(self, dialect) -> str:
+        operand_sql = self.operand.to_sql(dialect)
+        return f"({self.op} {operand_sql})"
+
+    def __repr__(self) -> str:
+        return f"({self.op} {self.operand})"
+
+
+class SQLComparison(SQLExpression):
+    """比较运算表达式"""
+
+    def __init__(self, left: SQLExpression, op: str, right: SQLExpression):
+        self.left = left
+        self.op = op
+        self.right = right
+
+    def to_sql(self, dialect) -> str:
+        left_sql = self.left.to_sql(dialect)
+        right_sql = self.right.to_sql(dialect)
+        return f"({left_sql} {self.op} {right_sql})"
+
+    def __repr__(self) -> str:
+        return f"({self.left} {self.op} {self.right})"
+
+
+class SQLLogicalOp(SQLExpression):
+    """逻辑运算表达式"""
+
+    def __init__(self, op: str, *operands: SQLExpression):
+        self.op = op
+        self.operands = operands
+
+    def to_sql(self, dialect) -> str:
+        if self.op == "NOT":
+            operand_sql = self.operands[0].to_sql(dialect)
+            return f"(NOT {operand_sql})"
+        parts = [op.to_sql(dialect) for op in self.operands]
+        return "(" + self.op + " " + " ".join(parts) + ")"
+
+    def __repr__(self) -> str:
+        if self.op == "NOT":
+            return f"(NOT {self.operands[0]})"
+        return f"({' ' + self.op + ' '.join(str(op) for op in self.operands)})"
+
+
+class SQLFunction(SQLExpression):
+    """SQL 函数调用表达式"""
+
+    def __init__(self, name: str, *args: SQLExpression, **kwargs: SQLExpression):
+        self.name = name
+        self.args = args
+        self.kwargs = kwargs
+
+    def to_sql(self, dialect) -> str:
+        args_sql = [arg.to_sql(dialect) for arg in self.args]
+        kwargs_sql = [f"{k}={v.to_sql(dialect)}" for k, v in self.kwargs.items()]
+        all_args = ", ".join(args_sql + kwargs_sql)
+        return f"{self.name}({all_args})"
+
+    def __repr__(self) -> str:
+        args_str = ", ".join(repr(a) for a in self.args)
+        return f"{self.name}({args_str})"
+
+
+class SQLCase(SQLExpression):
+    """CASE WHEN 表达式"""
+
+    def __init__(
+        self,
+        when_clauses: List[Tuple[SQLExpression, SQLExpression]],
+        else_: Optional[SQLExpression] = None,
+    ):
+        self.when_clauses = when_clauses
+        self.else_ = else_
+
+    def to_sql(self, dialect) -> str:
+        when_sql = [(cond.to_sql(dialect), val.to_sql(dialect)) for cond, val in self.when_clauses]
+        else_sql = self.else_.to_sql(dialect) if self.else_ else None
+        return dialect.func_case(when_sql, else_sql)
+
+    def __repr__(self) -> str:
+        else_str = f" ELSE {self.else_}" if self.else_ else ""
+        when_str = " ".join(f"WHEN {cond} THEN {val}" for cond, val in self.when_clauses)
+        return f"CASE {when_str}{else_str} END"
+
+
+def wrap_sql_expr(value: Any) -> SQLExpression:
+    """将值包装为 SQL 表达式"""
+    if isinstance(value, SQLExpression):
+        return value
+    if isinstance(value, str):
+        return ColumnRef(value)
+    return LiteralValue(value)
+
+
+def col(name: str, table: Optional[str] = None) -> ColumnRef:
+    """创建列引用"""
+    return ColumnRef(name, table)
+
+
+def lit(value: Any) -> LiteralValue:
+    """创建字面量"""
+    return LiteralValue(value)

QuantNodes/symbolic/functions.py

@@ -0,0 +1,433 @@
+# coding=utf-8
+"""
+符号计算引擎 - 技术指标函数库
+
+提供常用的技术分析指标函数。
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from QuantNodes.symbolic.expression import SQLExpression
+
+
+class TechnicalFunctions:
+    """
+    技术分析函数库
+
+    提供常用的技术指标函数，如移动平均、RSI、MACD 等。
+    这些函数生成 SQL 表达式，可在不同数据库方言中执行。
+    """
+
+    @staticmethod
+    def sma(expr: "SQLExpression", window: int) -> "SQLExpression":
+        """
+        简单移动平均
+
+        Args:
+            expr: 表达式
+            window: 窗口大小
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("avg", expr)
+
+    @staticmethod
+    def ema(expr: "SQLExpression", window: int) -> "SQLExpression":
+        """
+        指数移动平均
+
+        Args:
+            expr: 表达式
+            window: 窗口大小
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("expAverage", expr)
+
+    @staticmethod
+    def sum(expr: "SQLExpression", window: int) -> "SQLExpression":
+        """
+        滑动窗口求和
+
+        Args:
+            expr: 表达式
+            window: 窗口大小
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("sum", expr)
+
+    @staticmethod
+    def stddev(expr: "SQLExpression", window: int) -> "SQLExpression":
+        """
+        滑动窗口标准差
+
+        Args:
+            expr: 表达式
+            window: 窗口大小
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("stddevPop", expr)
+
+    @staticmethod
+    def variance(expr: "SQLExpression", window: int) -> "SQLExpression":
+        """
+        滑动窗口方差
+
+        Args:
+            expr: 表达式
+            window: 窗口大小
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("varPop", expr)
+
+    @staticmethod
+    def min(expr: "SQLExpression", window: int) -> "SQLExpression":
+        """
+        滑动窗口最小值
+
+        Args:
+            expr: 表达式
+            window: 窗口大小
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("min", expr)
+
+    @staticmethod
+    def max(expr: "SQLExpression", window: int) -> "SQLExpression":
+        """
+        滑动窗口最大值
+
+        Args:
+            expr: 表达式
+            window: 窗口大小
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("max", expr)
+
+    @staticmethod
+    def median(expr: "SQLExpression", window: int) -> "SQLExpression":
+        """
+        滑动窗口中位数
+
+        Args:
+            expr: 表达式
+            window: 窗口大小
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("median", expr)
+
+    @staticmethod
+    def quantile(expr: "SQLExpression", window: int, q: float) -> "SQLExpression":
+        """
+        滑动窗口分位数
+
+        Args:
+            expr: 表达式
+            window: 窗口大小
+            q: 分位数 (0-1)
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("quantile", expr)
+
+    @staticmethod
+    def delay(expr: "SQLExpression", n: int = 1) -> "SQLExpression":
+        """
+        滞后操作 (N日前值)
+
+        Args:
+            expr: 表达式
+            n: 滞后周期数
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("lagInFrame", expr)
+
+    @staticmethod
+    def delta(expr: "SQLExpression", n: int = 1) -> "SQLExpression":
+        """
+        差分 (与N日前的差值)
+
+        Args:
+            expr: 表达式
+            n: 滞后周期数
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        current = expr
+        lagged = SQLFunction("lagInFrame", expr)
+        return current - lagged
+
+    @staticmethod
+    def pct_change(expr: "SQLExpression", n: int = 1) -> "SQLExpression":
+        """
+        百分比变化
+
+        Args:
+            expr: 表达式
+            n: 滞后周期数
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        current = expr
+        lagged = SQLFunction("lagInFrame", expr)
+        return (current - lagged) / lagged
+
+    @staticmethod
+    def correlation(x: "SQLExpression", y: "SQLExpression", window: int) -> "SQLExpression":
+        """
+        滑动窗口相关系数
+
+        Args:
+            x: X 表达式
+            y: Y 表达式
+            window: 窗口大小
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("corr", x, y)
+
+    @staticmethod
+    def covariance(x: "SQLExpression", y: "SQLExpression", window: int) -> "SQLExpression":
+        """
+        滑动窗口协方差
+
+        Args:
+            x: X 表达式
+            y: Y 表达式
+            window: 窗口大小
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("covarPop", x, y)
+
+    @staticmethod
+    def rank(expr: "SQLExpression") -> "SQLExpression":
+        """
+        排名
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("rank")
+
+    @staticmethod
+    def dense_rank(expr: "SQLExpression") -> "SQLExpression":
+        """
+        密集排名
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("dense_rank")
+
+    @staticmethod
+    def row_number(expr: "SQLExpression") -> "SQLExpression":
+        """
+        行号
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("row_number")
+
+    @staticmethod
+    def cumsum(expr: "SQLExpression") -> "SQLExpression":
+        """
+        累积求和
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("sum", expr)
+
+    @staticmethod
+    def cumprod(expr: "SQLExpression") -> "SQLExpression":
+        """
+        累积求积
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("groupArray", expr)
+
+    @staticmethod
+    def running_diff(expr: "SQLExpression") -> "SQLExpression":
+        """
+        行间差分
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("runningDifference", expr)
+
+    @staticmethod
+    def abs(expr: "SQLExpression") -> "SQLExpression":
+        """
+        绝对值
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("abs", expr)
+
+    @staticmethod
+    def sqrt(expr: "SQLExpression") -> "SQLExpression":
+        """
+        平方根
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("sqrt", expr)
+
+    @staticmethod
+    def log(expr: "SQLExpression") -> "SQLExpression":
+        """
+        自然对数
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("log", expr)
+
+    @staticmethod
+    def log10(expr: "SQLExpression") -> "SQLExpression":
+        """
+        底数为10的对数
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("log10", expr)
+
+    @staticmethod
+    def pow(expr: "SQLExpression", power: float) -> "SQLExpression":
+        """
+        幂函数
+
+        Args:
+            expr: 表达式
+            power: 幂次
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("pow", expr)
+
+    @staticmethod
+    def round(expr: "SQLExpression", decimals: int = 0) -> "SQLExpression":
+        """
+        四舍五入
+
+        Args:
+            expr: 表达式
+            decimals: 小数位数
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("round", expr)
+
+    @staticmethod
+    def ceil(expr: "SQLExpression") -> "SQLExpression":
+        """
+        向上取整
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("ceil", expr)
+
+    @staticmethod
+    def floor(expr: "SQLExpression") -> "SQLExpression":
+        """
+        向下取整
+
+        Returns:
+            SQLExpression
+        """
+        from QuantNodes.symbolic.expression import SQLFunction
+        return SQLFunction("floor", expr)
+
+
+TA_FUNCTIONS = {
+    "sma": TechnicalFunctions.sma,
+    "ema": TechnicalFunctions.ema,
+    "sum": TechnicalFunctions.sum,
+    "stddev": TechnicalFunctions.stddev,
+    "variance": TechnicalFunctions.variance,
+    "min": TechnicalFunctions.min,
+    "max": TechnicalFunctions.max,
+    "median": TechnicalFunctions.median,
+    "quantile": TechnicalFunctions.quantile,
+    "delay": TechnicalFunctions.delay,
+    "delta": TechnicalFunctions.delta,
+    "pct_change": TechnicalFunctions.pct_change,
+    "correlation": TechnicalFunctions.correlation,
+    "covariance": TechnicalFunctions.covariance,
+    "rank": TechnicalFunctions.rank,
+    "dense_rank": TechnicalFunctions.dense_rank,
+    "row_number": TechnicalFunctions.row_number,
+    "cumsum": TechnicalFunctions.cumsum,
+    "cumprod": TechnicalFunctions.cumprod,
+    "running_diff": TechnicalFunctions.running_diff,
+    "abs": TechnicalFunctions.abs,
+    "sqrt": TechnicalFunctions.sqrt,
+    "log": TechnicalFunctions.log,
+    "log10": TechnicalFunctions.log10,
+    "pow": TechnicalFunctions.pow,
+    "round": TechnicalFunctions.round,
+    "ceil": TechnicalFunctions.ceil,
+    "floor": TechnicalFunctions.floor,
+}