data_pipeline/build_dataset.py

"""Build qlib DatasetH for model training workflow."""

from __future__ import annotations

from typing import Any

from qlib.contrib.data.handler import Alpha158
from qlib.data.dataset import DatasetH


def build_qlib_dataset(
    instruments: str = "csi300",
    start_time: str = "2010-01-01",
    end_time: str = "2025-12-31",
    fit_start: str = "2010-01-01",
    fit_end: str = "2021-12-31",
    segments: dict[str, tuple[str, str]] | None = None,
) -> DatasetH:
    """
    Build a standard Alpha158 dataset on top of qlib data handlers.
    This is the bridge between GP-mined factors and qlib model/backtest workflow.
    """
    if segments is None:
        segments = {
            "train": ("2010-01-01", "2021-12-31"),
            "valid": ("2022-01-01", "2023-12-31"),
            "test": ("2024-01-01", "2025-12-31"),
        }

    handler_kwargs: dict[str, Any] = {
        "start_time": start_time,
        "end_time": end_time,
        "fit_start_time": fit_start,
        "fit_end_time": fit_end,
        "instruments": instruments,
        "infer_processors": [],
        "learn_processors": [],
    }

    handler = Alpha158(**handler_kwargs)
    dataset = DatasetH(handler=handler, segments=segments)
    return dataset