ProStock/src/factors/engine/data_spec.py

"""数据规格和执行计划定义。

定义因子计算所需的数据规格和执行计划结构。
"""

from dataclasses import dataclass, field
from typing import Any, Dict, List, Literal, Optional, Set, Union

import polars as pl


@dataclass
class DataSpec:
    """数据规格定义（支持多表类型）。

    描述因子计算所需的数据表和字段，支持多种拼接类型。

    Attributes:
        table: 数据表名称
        columns: 需要的字段列表
        join_type: 拼接类型
            - "standard": 标准等值匹配（默认）
            - "asof_backward": 向后寻找最近历史数据（财务数据用）
        left_on: 左表 join 键（asof 模式下必须指定）
        right_on: 右表 join 键（asof 模式下必须指定）
    """

    table: str
    columns: List[str]
    join_type: Literal["standard", "asof_backward"] = "standard"
    left_on: Optional[str] = None  # 行情表日期列名
    right_on: Optional[str] = None  # 财务表日期列名

    def __post_init__(self):
        """验证 asof_backward 模式的参数。"""
        if self.join_type == "asof_backward":
            if not self.left_on or not self.right_on:
                raise ValueError("asof_backward 模式必须指定 left_on 和 right_on")


@dataclass
class ExecutionPlan:
    """执行计划。

    包含完整的执行所需信息：数据源、转换逻辑、输出格式。

    Attributes:
        data_specs: 数据规格列表
        polars_expr: Polars 表达式
        dependencies: 依赖的原始字段
        output_name: 输出因子名称
        factor_dependencies: 依赖的其他因子名称（用于分步执行）
    """

    data_specs: List[DataSpec]
    polars_expr: pl.Expr
    dependencies: Set[str]
    output_name: str
    factor_dependencies: Set[str] = field(default_factory=set)