"""指数数据同步服务。""" from __future__ import annotations import logging import gc from datetime import datetime, timedelta import polars as pl from app.indicators.pipeline import compute_enriched from app.services import kline_sync, preferences from app.tickflow.capabilities import Cap, CapabilitySet from app.tickflow.client import get_client from app.tickflow.repository import KlineRepository logger = logging.getLogger(__name__) def _quotes_to_index_instruments(resp) -> pl.DataFrame: """将数据源 quotes 响应规范为指数 instruments。""" if resp is None: return pl.DataFrame() if isinstance(resp, pl.DataFrame): df = resp elif hasattr(resp, "columns"): df = pl.from_pandas(resp.reset_index() if hasattr(resp, "reset_index") else resp) else: rows: list[dict] = [] for q in resp or []: item = q if isinstance(q, dict) else {} ext = item.get("ext") or {} symbol = item.get("symbol") if not symbol: continue rows.append({ "symbol": str(symbol), "name": ext.get("name") or item.get("name") or str(symbol), }) df = pl.DataFrame(rows) if df.is_empty() or "symbol" not in df.columns: return pl.DataFrame() rename = {"ts_code": "symbol"} df = df.rename({k: v for k, v in rename.items() if k in df.columns}) if "name" not in df.columns: if "ext" in df.columns: df = df.with_columns(pl.col("symbol").cast(pl.Utf8).alias("name")) else: df = df.with_columns(pl.col("symbol").cast(pl.Utf8).alias("name")) result = df.select([ pl.col("symbol").cast(pl.Utf8), pl.col("name").cast(pl.Utf8), ]).with_columns([ pl.col("symbol").str.split(".").list.first().alias("code"), pl.lit("index").alias("asset_type"), ]) return result.unique(subset=["symbol"], keep="last").sort("symbol") def sync_index_instruments(repo: KlineRepository) -> int: """同步 CN_Index 指数标的维表,返回指数数量。""" tf = get_client() resp = None errors: list[str] = [] for kwargs in ( {"universes": ["CN_Index"]}, {"universes": ["CN_Index"], "as_dataframe": False}, ): try: resp = tf.quotes.get_by_universes(**kwargs) if resp is not None and len(resp) > 0: break except Exception as e: # noqa: BLE001 errors.append(str(e)) resp = None if resp is None or len(resp) == 0: logger.warning("CN_Index universe returned empty: %s", "; ".join(errors)) return 0 instruments = _quotes_to_index_instruments(resp) if instruments.is_empty(): return 0 repo.save_index_instruments(instruments) repo.refresh_index_views() return instruments.height def sync_and_persist_index_daily( repo: KlineRepository, capset: CapabilitySet, count: int | None = None, start_date: datetime | None = None, end_date: datetime | None = None, ) -> int: """同步指数日K到独立 parquet,并计算指数 enriched。""" if not capset.has(Cap.KLINE_DAILY_BATCH): return 0 instruments = repo.get_index_instruments() if instruments.is_empty(): sync_index_instruments(repo) instruments = repo.get_index_instruments() if instruments.is_empty() or "symbol" not in instruments.columns: return 0 symbols = sorted(set(instruments["symbol"].to_list())) lim = capset.limits(Cap.KLINE_DAILY_BATCH) batch_size = preferences.get_index_daily_batch_size() if lim and lim.batch: batch_size = min(batch_size, lim.batch) rpm = lim.rpm if lim else None end_time = end_date or datetime.now() start_time = start_date or (end_time - timedelta(days=365)) total_rows = 0 interval = (60.0 / rpm) if rpm else 0 chunks = [symbols[i:i + batch_size] for i in range(0, len(symbols), batch_size)] for i, chunk in enumerate(chunks): if i > 0 and interval > 0 and len(chunks) > rpm: import time time.sleep(interval) raw = kline_sync.sync_daily_batch( chunk, count=count, batch_size=None, start_time=start_time, end_time=end_time, ) if raw.is_empty(): continue repo.append_index_daily(raw) enriched = compute_enriched(raw, factors=None, instruments=None) repo.append_index_enriched(enriched) total_rows += raw.height logger.info("index daily synced: %d/%d chunks, +%d rows", i + 1, len(chunks), raw.height) del raw, enriched gc.collect() repo.refresh_index_views() return total_rows