quantnodes 3.0.0__py3-none-any.whl

QuantNodes/core/_lookback_helpers.py

@@ -0,0 +1,49 @@
+# -*- coding: utf-8 -*-
+"""LookBack 窗口计算的共用辅助函数"""
+import numpy as np
+
+
+def compute_lookback_params(lookback, lookback_mode):
+    """计算 LookBack 窗口参数
+
+    Args:
+        lookback: list of lookback values for descriptors
+        lookback_mode: list of lookback modes for descriptors
+
+    Returns:
+        tuple: (StartIndAndLen, MaxLookBack, MaxLen)
+            - StartIndAndLen: list of (start_ind, length) tuples
+            - MaxLookBack: maximum lookback value
+            - MaxLen: maximum window length
+    """
+    StartIndAndLen, MaxLookBack, MaxLen = [], 0, 1
+    for i, iLookBack in enumerate(lookback):
+        if lookback_mode[i] == "滚动窗口":
+            StartIndAndLen.append((iLookBack, iLookBack + 1))
+            MaxLen = max(MaxLen, iLookBack + 1)
+        else:
+            StartIndAndLen.append((iLookBack, np.inf))
+            MaxLen = np.inf
+        MaxLookBack = max(MaxLookBack, iLookBack)
+    return StartIndAndLen, MaxLookBack, MaxLen
+
+
+def extend_dt_ruler(dt_ruler, dts, max_lookback):
+    """扩展时间标尺向前回溯
+
+    如果 start_ind >= max_lookback，直接返回 dt_ruler[start_ind-max_lookback:]
+    否则在前面补 None 凑够 max_lookback 个。
+
+    Args:
+        dt_ruler: 时间标尺
+        dts: 当前时间序列
+        max_lookback: 最大回溯值
+
+    Returns:
+        list: 扩展后的时间标尺
+    """
+    start_ind = dt_ruler.index(dts[0]) if dts[0] in dt_ruler else 0
+    if start_ind >= max_lookback:
+        return list(dt_ruler[start_ind - max_lookback:])
+    else:
+        return [None] * (max_lookback - start_ind) + list(dt_ruler)

QuantNodes/core/ast_parser.py

@@ -0,0 +1,198 @@
+# coding=utf-8
+"""
+AST 安全表达式解析器
+
+基于 Python 标准库的 ast 模块，安全解析表达式字符串，
+不使用 eval，避免代码注入风险。
+"""
+
+from __future__ import annotations
+
+import ast
+import logging
+
+from QuantNodes.core.expression import (
+    Expression,
+    InputExpr,
+    ConstantExpr,
+    VariableExpr,
+    AttributeExpr,
+    SubscriptExpr,
+    MethodCallExpr,
+    BinaryOpExpr,
+    UnaryOpExpr,
+    ComparisonExpr,
+    LogicalOpExpr,
+    ALLOWED_AST_NODES,
+    FORBIDDEN_METHODS,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+def parse_expression(expr_str: str) -> Expression:
+    """
+    安全解析字符串表达式
+
+    Args:
+        expr_str: 表达式字符串，如 "df['close'] > 50"
+
+    Returns:
+        Expression 对象
+
+    Raises:
+        SyntaxError: 表达式语法错误
+        ValueError: 包含不支持的操作
+    """
+    try:
+        tree = ast.parse(expr_str, mode="eval")
+    except SyntaxError as e:
+        raise SyntaxError(f"Invalid expression syntax: {e}") from e
+
+    # 安全检查
+    _validate_ast(tree)
+
+    # 转换为表达式对象
+    return _ast_to_expr(tree.body)
+
+
+def _validate_ast(node: ast.AST) -> None:
+    """递归验证 AST 节点是否在白名单中"""
+    if type(node) not in ALLOWED_AST_NODES:
+        raise ValueError(f"Unsupported AST node type: {type(node).__name__}")
+
+    for child in ast.iter_child_nodes(node):
+        _validate_ast(child)
+
+
+def _ast_to_expr(node: ast.AST) -> Expression:
+    """将 AST 节点转换为 Expression 对象"""
+
+    # 常量
+    if isinstance(node, ast.Constant):
+        return ConstantExpr(node.value)
+
+    # 变量名 - 假设是输入数据的属性或列名
+    if isinstance(node, ast.Name):
+        # 特殊处理常见的输入变量名
+        if node.id in ("df", "input", "data", "x", "result"):
+            return InputExpr()
+        # 其他名称视为属性访问
+        return VariableExpr(node.id)
+
+    # 属性访问
+    if isinstance(node, ast.Attribute):
+        expr = _ast_to_expr(node.value)
+        return AttributeExpr(expr, node.attr)
+
+    # 下标访问
+    if isinstance(node, ast.Subscript):
+        expr = _ast_to_expr(node.value)
+        key = _ast_to_expr(node.slice)
+        return SubscriptExpr(expr, key)
+
+    # 方法调用
+    if isinstance(node, ast.Call):
+        if isinstance(node.func, ast.Attribute):
+            expr = _ast_to_expr(node.func.value)
+            method_name = node.func.attr
+
+            if method_name in FORBIDDEN_METHODS:
+                raise ValueError(f"Forbidden method: {method_name}")
+
+            args = tuple(_ast_to_expr(arg) for arg in node.args)
+            kwargs = {
+                kw.arg: _ast_to_expr(kw.value)
+                for kw in node.keywords
+                if kw.arg is not None
+            }
+            return MethodCallExpr(expr, method_name, args, kwargs)
+        raise ValueError("Unsupported call type")
+
+    # 二元运算
+    if isinstance(node, ast.BinOp):
+        left = _ast_to_expr(node.left)
+        right = _ast_to_expr(node.right)
+        op = _bin_op_to_str(node.op)
+        return BinaryOpExpr(left, op, right)
+
+    # 一元运算
+    if isinstance(node, ast.UnaryOp):
+        operand = _ast_to_expr(node.operand)
+        op = _unary_op_to_str(node.op)
+        return UnaryOpExpr(op, operand)
+
+    # 比较运算
+    if isinstance(node, ast.Compare):
+        left = _ast_to_expr(node.left)
+        # 简化：只处理单个比较
+        op = _cmp_op_to_str(node.ops[0])
+        right = _ast_to_expr(node.comparators[0])
+        return ComparisonExpr(left, op, right)
+
+    # 逻辑运算
+    if isinstance(node, ast.BoolOp):
+        op = _bool_op_to_str(node.op)
+        operands = tuple(_ast_to_expr(v) for v in node.values)
+        return LogicalOpExpr(op, *operands)
+
+    raise ValueError(f"Unsupported AST node: {type(node).__name__}")
+
+
+def _bin_op_to_str(op: ast.operator) -> str:
+    """将二元运算 AST 节点转换为运算符字符串"""
+    mapping = {
+        ast.Add: "+",
+        ast.Sub: "-",
+        ast.Mult: "*",
+        ast.Div: "/",
+        ast.FloorDiv: "//",
+        ast.Mod: "%",
+        ast.Pow: "**",
+    }
+    result = mapping.get(type(op))
+    if result is None:
+        raise ValueError(f"Unsupported binary operator: {type(op).__name__}")
+    return result
+
+
+def _unary_op_to_str(op: ast.unaryop) -> str:
+    """将一元运算 AST 节点转换为运算符字符串"""
+    mapping = {
+        ast.UAdd: "+",
+        ast.USub: "-",
+        ast.Not: "not",
+    }
+    result = mapping.get(type(op))
+    if result is None:
+        raise ValueError(f"Unsupported unary operator: {type(op).__name__}")
+    return result
+
+
+def _cmp_op_to_str(op: ast.cmpop) -> str:
+    """将比较运算 AST 节点转换为运算符字符串"""
+    mapping = {
+        ast.Gt: ">",
+        ast.GtE: ">=",
+        ast.Lt: "<",
+        ast.LtE: "<=",
+        ast.Eq: "==",
+        ast.NotEq: "!=",
+    }
+    result = mapping.get(type(op))
+    if result is None:
+        raise ValueError(f"Unsupported comparison operator: {type(op).__name__}")
+    return result
+
+
+def _bool_op_to_str(op: ast.boolop) -> str:
+    """将逻辑运算 AST 节点转换为运算符字符串"""
+    mapping = {
+        ast.And: "and",
+        ast.Or: "or",
+    }
+    result = mapping.get(type(op))
+    if result is None:
+        raise ValueError(f"Unsupported logical operator: {type(op).__name__}")
+    return result

QuantNodes/core/base.py

@@ -0,0 +1,61 @@
+# coding=utf-8
+
+from typing import Any, Dict, Optional
+from pydantic import BaseModel as PydanticBaseModel, ConfigDict
+
+
+class BaseModel(PydanticBaseModel):
+    """基础模型类"""
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="forbid",
+        frozen=False,
+    )
+
+
+class QuantNodesBase:
+    """所有业务类的基类"""
+
+    def __init__(self, name: Optional[str] = None):
+        self.name = name or self.__class__.__name__
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}: {self.name}>"
+
+
+class QuantNodesError(Exception):
+    """基础异常类"""
+    code = "QUANTNODES_ERROR"
+
+    def __init__(self, message: str, details: Optional[Dict[str, Any]] = None):
+        super().__init__(message)
+        self.message = message
+        self.details = details or {}
+
+    def __str__(self) -> str:
+        return f"[{self.code}] {self.message}"
+
+
+class ConfigError(QuantNodesError):
+    """配置异常"""
+    code = "CONFIG_ERROR"
+
+
+class DatabaseError(QuantNodesError):
+    """数据库异常"""
+    code = "DATABASE_ERROR"
+
+
+class FactorError(QuantNodesError):
+    """因子异常"""
+    code = "FACTOR_ERROR"
+
+
+class BacktestError(QuantNodesError):
+    """回测异常"""
+    code = "BACKTEST_ERROR"
+
+
+class ValidationError(QuantNodesError):
+    """校验异常"""
+    code = "VALIDATION_ERROR"

QuantNodes/core/cache_manager.py

@@ -0,0 +1,344 @@
+# coding=utf-8
+"""
+缓存管理器
+
+替代 QuantStudio 中与缓存和进程管理相关的功能
+"""
+
+import gc
+import mmap
+import os
+import pickle
+import shelve
+from multiprocessing import Lock
+from typing import Any, Dict, List, Optional
+
+import pandas as pd
+
+from QuantNodes.core.tools import get_shelve_file_suffix
+
+
+class ErgodicMode:
+    """
+    遍历模式参数对象
+
+    管理因子遍历模式下的缓存策略
+    """
+
+    def __init__(
+        self,
+        forward_period: int = 600,
+        backward_period: int = 1,
+        cache_mode: str = "因子",
+        max_factor_cache_num: int = 60,
+        max_id_cache_num: int = 10000,
+        cache_size: int = 300,
+        ergodic_dts: Optional[List] = None,
+        ergodic_ids: Optional[List] = None,
+    ):
+        self.ForwardPeriod = forward_period
+        self.BackwardPeriod = backward_period
+        self.CacheMode = cache_mode
+        self.MaxFactorCacheNum = max_factor_cache_num
+        self.MaxIDCacheNum = max_id_cache_num
+        self.CacheSize = cache_size
+        self.ErgodicDTs = ergodic_dts or []
+        self.ErgodicIDs = ergodic_ids or []
+        self._isStarted = False
+        self._CurDT = None
+        self._DateTimes = None
+        self._IDs = None
+        self._CurInd = None
+        self._DTNum = None
+        self._CacheDTs = None
+        self._CacheData = None
+        self._CacheFactorNum = None
+        self._CacheIDNum = None
+        self._FactorReadNum = None
+        self._IDReadNum = None
+        self._Queue2SubProcess = None
+        self._Queue2MainProcess = None
+        self._TagName = None
+        self._MMAPCacheData = None
+        self._CacheDataProcess = None
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        if "_CacheDataProcess" in state:
+            state["_CacheDataProcess"] = None
+        return state
+
+
+class OperationMode:
+    """
+    运算模式参数对象
+
+    管理因子运算模式下的多进程调度
+    """
+
+    def __init__(
+        self,
+        ft: Any = None,
+        factor_names: Optional[List[str]] = None,
+        subprocess_num: int = 0,
+    ):
+        self._FT = ft
+        self._isStarted = False
+        self._Factors = []
+        self._FactorDict = {}
+        self._FactorID = {}
+        self._FactorStartDT = {}
+        self._FactorPrepareIDs = {}
+        self._iPID = "0"
+        self._PIDs = []
+        self._PID_IDs = {}
+        self._PID_Lock = {}
+        self._CacheDir = None
+        self._RawDataDir = ""
+        self._CacheDataDir = ""
+        self._Event = {}
+        self._DateTimes = []
+        self._IDs = []
+        self._DTRuler = []
+        self.SubProcessNum = subprocess_num
+        self.FactorNames = factor_names or []
+        self.FileSuffix = get_shelve_file_suffix()
+        if self.FileSuffix:
+            self.FileSuffix = "." + self.FileSuffix
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        if self._CacheDir is not None:
+            state["_CacheDir"] = self._CacheDir.name
+        return state
+
+
+def prepare_mmap_factor_cache_data(ft: Any, mmap_cache: Any) -> int:
+    """
+    基于 mmap 的因子缓冲数据准备子进程
+    """
+    CacheData = {}
+    CacheDTs = []
+    CacheSize = int(ft.ErgodicMode.CacheSize * 2**20)
+
+    if os.name == "nt":
+        MMAPCacheData = mmap.mmap(-1, CacheSize, tagname=ft.ErgodicMode._TagName)
+    else:
+        MMAPCacheData = mmap_cache
+
+    while True:
+        Task = ft.ErgodicMode._Queue2SubProcess.get()
+        if Task is None:
+            break
+
+        if Task[0] is None and Task[1] is None:
+            CacheDataByte = pickle.dumps(CacheData)
+            DataLen = len(CacheDataByte)
+            for i in range(int(DataLen / CacheSize) + 1):
+                iStartInd = i * CacheSize
+                iEndInd = min((i + 1) * CacheSize, DataLen)
+                if iEndInd > iStartInd:
+                    MMAPCacheData.seek(0)
+                    MMAPCacheData.write(CacheDataByte[iStartInd:iEndInd])
+                    ft.ErgodicMode._Queue2MainProcess.put(iEndInd - iStartInd)
+                    ft.ErgodicMode._Queue2SubProcess.get()
+            ft.ErgodicMode._Queue2MainProcess.put(0)
+            del CacheDataByte
+            gc.collect()
+        elif Task[0] is None:
+            NewFactors, PopFactors = Task[1]
+            for iFactorName in PopFactors:
+                CacheData.pop(iFactorName)
+            if NewFactors:
+                if CacheDTs:
+                    CacheData.update(
+                        dict(
+                            ft.__QN_calc_data__(
+                                raw_data=ft.__QN_prepare_raw_data__(
+                                    factor_names=NewFactors,
+                                    ids=ft.ErgodicMode._IDs,
+                                    dts=CacheDTs,
+                                ),
+                                factor_names=NewFactors,
+                                ids=ft.ErgodicMode._IDs,
+                                dts=CacheDTs,
+                            )
+                        )
+                    )
+                else:
+                    CacheData.update(
+                        {
+                            iFactorName: pd.DataFrame(index=CacheDTs, columns=ft.ErgodicMode._IDs)
+                            for iFactorName in NewFactors
+                        }
+                    )
+        else:
+            CurInd = Task[0] + ft.ErgodicMode.ForwardPeriod + 1
+            DTNum = len(ft.ErgodicMode._DateTimes)
+            if CurInd < DTNum:
+                OldCacheDTs = set(CacheDTs)
+                CacheDTs = ft.ErgodicMode._DateTimes[
+                    max((0, CurInd - ft.ErgodicMode.BackwardPeriod)) : min(
+                        (DTNum, CurInd + ft.ErgodicMode.ForwardPeriod + 1)
+                    )
+                ].tolist()
+                NewCacheDTs = sorted(set(CacheDTs).difference(OldCacheDTs))
+                if CacheData:
+                    isDisjoint = OldCacheDTs.isdisjoint(CacheDTs)
+                    CacheFactorNames = list(CacheData.keys())
+                    if NewCacheDTs:
+                        NewCacheData = ft.__QN_calc_data__(
+                            raw_data=ft.__QN_prepare_raw_data__(
+                                factor_names=CacheFactorNames,
+                                ids=ft.ErgodicMode._IDs,
+                                dts=NewCacheDTs,
+                            ),
+                            factor_names=CacheFactorNames,
+                            ids=ft.ErgodicMode._IDs,
+                            dts=NewCacheDTs,
+                        )
+                    else:
+                        NewCacheData = pd.DataFrame(
+                            index=CacheFactorNames, columns=CacheDTs
+                        )
+                    for iFactorName in CacheData:
+                        if isDisjoint:
+                            CacheData[iFactorName] = NewCacheData[iFactorName]
+                        else:
+                            CacheData[iFactorName] = CacheData[iFactorName].loc[CacheDTs, :]
+                            CacheData[iFactorName].loc[NewCacheDTs, :] = NewCacheData[iFactorName]
+                    NewCacheData = None
+    return 0
+
+
+def prepare_mmap_id_cache_data(ft: Any, mmap_cache: Any) -> int:
+    """
+    基于 mmap 的 ID 缓冲数据准备子进程
+    """
+    CacheData = {}
+    CacheDTs = []
+    CacheSize = int(ft.ErgodicMode.CacheSize * 2**20)
+
+    if os.name == "nt":
+        MMAPCacheData = mmap.mmap(-1, CacheSize, tagname=ft.ErgodicMode._TagName)
+    else:
+        MMAPCacheData = mmap_cache
+
+    while True:
+        Task = ft.ErgodicMode._Queue2SubProcess.get()
+        if Task is None:
+            break
+
+        if Task[0] is None and Task[1] is None:
+            CacheDataByte = pickle.dumps(CacheData)
+            DataLen = len(CacheDataByte)
+            for i in range(int(DataLen / CacheSize) + 1):
+                iStartInd = i * CacheSize
+                iEndInd = min((i + 1) * CacheSize, DataLen)
+                if iEndInd > iStartInd:
+                    MMAPCacheData.seek(0)
+                    MMAPCacheData.write(CacheDataByte[iStartInd:iEndInd])
+                    ft.ErgodicMode._Queue2MainProcess.put(iEndInd - iStartInd)
+                    ft.ErgodicMode._Queue2SubProcess.get()
+            ft.ErgodicMode._Queue2MainProcess.put(0)
+            del CacheDataByte
+            gc.collect()
+        elif Task[0] is None:
+            NewID, PopID = Task[1]
+            if PopID:
+                CacheData.pop(PopID)
+            if NewID:
+                if CacheDTs:
+                    CacheData[NewID] = ft.__QN_calc_data__(
+                        raw_data=ft.__QN_prepare_raw_data__(
+                            factor_names=ft.FactorNames,
+                            ids=[NewID],
+                            dts=CacheDTs,
+                        ),
+                        factor_names=ft.FactorNames,
+                        ids=[NewID],
+                        dts=CacheDTs,
+                    ).iloc[:, :, 0]
+                else:
+                    CacheData[NewID] = pd.DataFrame(index=CacheDTs, columns=ft.FactorNames)
+        else:
+            CurInd = Task[0] + ft.ErgodicMode.ForwardPeriod + 1
+            DTNum = len(ft.ErgodicMode._DateTimes)
+            if CurInd < DTNum:
+                OldCacheDTs = set(CacheDTs)
+                CacheDTs = ft.ErgodicMode._DateTimes[
+                    max((0, CurInd - ft.ErgodicMode.BackwardPeriod)) : min(
+                        (DTNum, CurInd + ft.ErgodicMode.ForwardPeriod + 1)
+                    )
+                ].tolist()
+                NewCacheDTs = sorted(set(CacheDTs).difference(OldCacheDTs))
+                if CacheData:
+                    isDisjoint = OldCacheDTs.isdisjoint(CacheDTs)
+                    CacheIDs = list(CacheData.keys())
+                    if NewCacheDTs:
+                        NewCacheData = ft.__QN_calc_data__(
+                            raw_data=ft.__QN_prepare_raw_data__(
+                                factor_names=ft.FactorNames,
+                                ids=CacheIDs,
+                                dts=NewCacheDTs,
+                            ),
+                            factor_names=ft.FactorNames,
+                            ids=CacheIDs,
+                            dts=NewCacheDTs,
+                        )
+                    else:
+                        NewCacheData = pd.DataFrame(
+                            index=ft.FactorNames, columns=CacheDTs
+                        )
+                    for iID in CacheData:
+                        if isDisjoint:
+                            CacheData[iID] = NewCacheData[iID]
+                        else:
+                            CacheData[iID] = CacheData[iID].loc[CacheDTs, :]
+                            CacheData[iID].loc[NewCacheDTs, :] = NewCacheData[iID]
+                    NewCacheData = None
+    return 0
+
+
+def save_raw_data(
+    raw_data: Any,
+    factor_names: List[str],
+    raw_data_dir: str,
+    pid_ids: Dict[str, List[str]],
+    file_name: str,
+    pid_lock: Dict[str, Lock],
+) -> int:
+    """
+    保存原始数据到文件
+    """
+    if raw_data is None:
+        return 0
+    if isinstance(raw_data, pd.DataFrame) and ("ID" in raw_data):
+        raw_data = raw_data.set_index(["ID"])
+        CommonCols = raw_data.columns.difference(factor_names).tolist()
+        AllIDs = set(raw_data.index)
+        for iPID, iIDs in pid_ids.items():
+            with shelve.open(raw_data_dir + os.sep + iPID + os.sep + file_name) as iFile:
+                iInterIDs = sorted(AllIDs.intersection(iIDs))
+                iData = raw_data.loc[iInterIDs]
+                if factor_names:
+                    for jFactorName in factor_names:
+                        iFile[jFactorName] = iData[CommonCols + [jFactorName]].reset_index()
+                else:
+                    iFile["RawData"] = iData[CommonCols].reset_index()
+                iFile["_QN_IDs"] = iIDs
+    else:
+        for iPID, iIDs in pid_ids.items():
+            with shelve.open(raw_data_dir + os.sep + iPID + os.sep + file_name) as iFile:
+                iFile["RawData"] = raw_data
+                iFile["_QN_IDs"] = iIDs
+    return 0
+
+
+__all__ = [
+    "ErgodicMode",
+    "OperationMode",
+    "prepare_mmap_factor_cache_data",
+    "prepare_mmap_id_cache_data",
+    "save_raw_data",
+]