ProStock/src/factors/engine/planner.py

"""执行计划生成器。

整合编译器和翻译器，生成完整的执行计划。
"""

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