refactor(factors): 拆分 engine.py 为模块化包
将单文件 engine.py (1064行) 拆分为 engine/ 包: - 数据规格、路由器、计划器、计算引擎、因子引擎分离 - 保持向后兼容,API 无变化
This commit is contained in:
170
src/factors/engine/planner.py
Normal file
170
src/factors/engine/planner.py
Normal file
@@ -0,0 +1,170 @@
|
||||
"""执行计划生成器。
|
||||
|
||||
整合编译器和翻译器,生成完整的执行计划。
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Set, Union
|
||||
|
||||
from src.factors.dsl import (
|
||||
Node,
|
||||
Symbol,
|
||||
FunctionNode,
|
||||
BinaryOpNode,
|
||||
UnaryOpNode,
|
||||
Constant,
|
||||
)
|
||||
from src.factors.compiler import DependencyExtractor
|
||||
from src.factors.translator import PolarsTranslator
|
||||
from src.factors.engine.data_spec import DataSpec, ExecutionPlan
|
||||
|
||||
|
||||
class ExecutionPlanner:
|
||||
"""执行计划生成器。
|
||||
|
||||
整合编译器和翻译器,生成完整的执行计划。
|
||||
|
||||
Attributes:
|
||||
compiler: 依赖提取器
|
||||
translator: Polars 翻译器
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""初始化执行计划生成器。"""
|
||||
self.compiler = DependencyExtractor()
|
||||
self.translator = PolarsTranslator()
|
||||
|
||||
def create_plan(
|
||||
self,
|
||||
expression: Node,
|
||||
output_name: str = "factor",
|
||||
data_specs: Optional[List[DataSpec]] = None,
|
||||
) -> ExecutionPlan:
|
||||
"""从表达式创建执行计划。
|
||||
|
||||
Args:
|
||||
expression: DSL 表达式节点
|
||||
output_name: 输出因子名称
|
||||
data_specs: 预定义的数据规格,None 时自动推导
|
||||
|
||||
Returns:
|
||||
执行计划对象
|
||||
"""
|
||||
# 1. 提取依赖
|
||||
dependencies = self.compiler.extract_dependencies(expression)
|
||||
|
||||
# 2. 翻译为 Polars 表达式
|
||||
polars_expr = self.translator.translate(expression)
|
||||
|
||||
# 3. 推导或验证数据规格
|
||||
if data_specs is None:
|
||||
data_specs = self._infer_data_specs(dependencies, expression)
|
||||
|
||||
return ExecutionPlan(
|
||||
data_specs=data_specs,
|
||||
polars_expr=polars_expr,
|
||||
dependencies=dependencies,
|
||||
output_name=output_name,
|
||||
)
|
||||
|
||||
def _infer_data_specs(
|
||||
self,
|
||||
dependencies: Set[str],
|
||||
expression: Node,
|
||||
) -> List[DataSpec]:
|
||||
"""从依赖推导数据规格。
|
||||
|
||||
根据表达式中的函数类型推断回看天数需求。
|
||||
基础行情字段(open, high, low, close, vol, amount, pre_close, change, pct_chg)
|
||||
默认从 pro_bar 表获取。
|
||||
|
||||
Args:
|
||||
dependencies: 依赖的字段集合
|
||||
expression: 表达式节点
|
||||
|
||||
Returns:
|
||||
数据规格列表
|
||||
"""
|
||||
# 计算最大回看窗口
|
||||
max_window = self._extract_max_window(expression)
|
||||
lookback_days = max(1, max_window)
|
||||
|
||||
# 基础行情字段集合(这些字段从 pro_bar 表获取)
|
||||
pro_bar_fields = {
|
||||
"open",
|
||||
"high",
|
||||
"low",
|
||||
"close",
|
||||
"vol",
|
||||
"amount",
|
||||
"pre_close",
|
||||
"change",
|
||||
"pct_chg",
|
||||
"turnover_rate",
|
||||
"volume_ratio",
|
||||
}
|
||||
|
||||
# 将依赖分为 pro_bar 字段和其他字段
|
||||
pro_bar_deps = dependencies & pro_bar_fields
|
||||
other_deps = dependencies - pro_bar_fields
|
||||
|
||||
data_specs = []
|
||||
|
||||
# pro_bar 表的数据规格
|
||||
if pro_bar_deps:
|
||||
data_specs.append(
|
||||
DataSpec(
|
||||
table="pro_bar",
|
||||
columns=sorted(pro_bar_deps),
|
||||
lookback_days=lookback_days,
|
||||
)
|
||||
)
|
||||
|
||||
# 其他字段从 daily 表获取
|
||||
if other_deps:
|
||||
data_specs.append(
|
||||
DataSpec(
|
||||
table="daily",
|
||||
columns=sorted(other_deps),
|
||||
lookback_days=lookback_days,
|
||||
)
|
||||
)
|
||||
|
||||
return data_specs
|
||||
|
||||
def _extract_max_window(self, node: Node) -> int:
|
||||
"""从表达式中提取最大窗口大小。
|
||||
|
||||
Args:
|
||||
node: AST 节点
|
||||
|
||||
Returns:
|
||||
最大窗口大小,无时序函数返回 1
|
||||
"""
|
||||
if isinstance(node, FunctionNode):
|
||||
window = 1
|
||||
# 检查函数参数中的窗口大小
|
||||
for arg in node.args:
|
||||
if (
|
||||
isinstance(arg, Constant)
|
||||
and isinstance(arg.value, int)
|
||||
and arg.value > window
|
||||
):
|
||||
window = arg.value
|
||||
|
||||
# 递归检查子表达式
|
||||
for arg in node.args:
|
||||
if isinstance(arg, Node) and not isinstance(arg, Constant):
|
||||
window = max(window, self._extract_max_window(arg))
|
||||
|
||||
return window
|
||||
|
||||
elif isinstance(node, BinaryOpNode):
|
||||
return max(
|
||||
self._extract_max_window(node.left),
|
||||
self._extract_max_window(node.right),
|
||||
)
|
||||
|
||||
elif isinstance(node, UnaryOpNode):
|
||||
return self._extract_max_window(node.operand)
|
||||
|
||||
return 1
|
||||
Reference in New Issue
Block a user