sanguo_vnpy_v2/sanguo_factor/alpha_lab.py

"""AlphaLab session management - lazy import vnpy.alpha to avoid ImportError when alphalens missing."""
import sys
import os
_VNPY_SRC = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "vnpy_v4.4.0"))
if _VNPY_SRC not in sys.path:
    sys.path.insert(0, _VNPY_SRC)

from collections import OrderedDict

# Maximum number of symbols to cache in _loaded_bars to prevent unbounded growth
# Single-user internal tool - 50 symbols is a generous ceiling
_MAX_CACHED_SYMBOLS = 50


class AlphaLabSession:
    """Session manager for vnpy.alpha AlphaLab operations."""

    def __init__(self, lab_path: str):
        """
        Initialize AlphaLab session.

        Args:
            lab_path: Path to AlphaLab directory
        """
        from vnpy.alpha.lab import AlphaLab

        self.lab_path = lab_path
        self.lab = AlphaLab(lab_path)
        self._loaded_symbols: list[str] = []
        # Use OrderedDict to maintain insertion order for LRU eviction
        self._loaded_bars: OrderedDict[str, list] = OrderedDict()
    
    def load_symbols(self, symbols: list[str], start: str, end: str, cfg) -> None:
        """
        Load symbol data from database and save to AlphaLab.

        Args:
            symbols: List of vt_symbols to load
            start: Start date (YYYY-MM-DD)
            end: End date (YYYY-MM-DD)
            cfg: Database configuration object
        """
        from sanguo_data.datareader import read_db_daily
        from .data_adapter import save_alpha_lab_data

        for symbol in symbols:
            bars = read_db_daily(symbol, start, end, cfg)
            if bars:
                save_alpha_lab_data(bars, self.lab_path)
                # Cache bars for compute_factors with LRU eviction
                if symbol not in self._loaded_symbols:
                    self._loaded_symbols.append(symbol)

                # Add or update symbol in cache (moves to end if exists)
                if symbol in self._loaded_bars:
                    del self._loaded_bars[symbol]  # Remove to re-insert for LRU
                self._loaded_bars[symbol] = bars

                # Enforce cache cap - evict oldest symbol if exceeded
                while len(self._loaded_bars) > _MAX_CACHED_SYMBOLS:
                    oldest_symbol = next(iter(self._loaded_bars))
                    del self._loaded_bars[oldest_symbol]
                    if oldest_symbol in self._loaded_symbols:
                        self._loaded_symbols.remove(oldest_symbol)

    def compute_factors(self, factor_names: list[str], train_period: tuple, valid_period: tuple, test_period: tuple):
        """
        Compute factors using cached bars and vnpy.alpha AlphaDataset.

        Args:
            factor_names: List of factor names to compute
            train_period: Training period tuple (start, end)
            valid_period: Validation period tuple (start, end)
            test_period: Test period tuple (start, end)

        Returns:
            polars DataFrame with computed factors for test period
        """
        # Lazy imports to avoid ImportError on local Python 3.14 without polars/vnpy.alpha
        import polars as pl
        from vnpy.alpha.dataset import AlphaDataset, Segment
        from .registry import get_factor
        from .data_adapter import convert_bars_to_alpha_df

        # Convert period boundaries to Asia/Shanghai-aware datetime objects
        # This ensures vnpy.alpha's to_datetime() preserves timezone awareness,
        # preventing SchemaError when comparing aware datetime column with naive literals
        from datetime import datetime
        from zoneinfo import ZoneInfo
        _SH = ZoneInfo("Asia/Shanghai")

        def _to_aware_period(period: tuple) -> tuple:
            """Convert period boundary strings or naive datetimes to Asia/Shanghai-aware datetimes."""
            start, end = period
            def conv(x):
                if isinstance(x, datetime):
                    return x if x.tzinfo else x.replace(tzinfo=_SH)
                return datetime.strptime(x, "%Y-%m-%d").replace(tzinfo=_SH)
            return (conv(start), conv(end))

        train_period = _to_aware_period(train_period)
        valid_period = _to_aware_period(valid_period)
        test_period = _to_aware_period(test_period)

        # Gather all cached bars across loaded symbols
        all_bars = []
        for symbol in self._loaded_symbols:
            all_bars.extend(self._loaded_bars.get(symbol, []))

        # Convert bars to AlphaLab DataFrame format
        df = convert_bars_to_alpha_df(all_bars)

        # Localize the datetime column to Asia/Shanghai-aware to match vnpy.alpha's expectations
        # This ensures the DataFrame's datetime column has the same timezone as the period boundaries
        df = df.with_columns(
            pl.col("datetime").dt.replace_time_zone("Asia/Shanghai")
        )

        # Create AlphaDataset with the specified periods
        ds = AlphaDataset(df, train_period, valid_period, test_period)

        # Add each factor to the dataset
        for name in factor_names:
            factor = get_factor(name)
            if factor is None:
                continue  # Skip unknown factors
            ds.add_feature(name, factor["expression"])

        # Prepare data (compute features)
        ds.prepare_data(max_workers=1)

        # Return test period data
        return ds.fetch_raw(Segment.TEST)