Merge pull request #207 from shuolol/feat/strategy-index-data

feat(strategy): 策略可读取指数/ETF日K (market_data 模块 + 白名单 + 测试)

backend/app/strategy/ai_generator.py

@@ -346,6 +346,7 @@ META = {{...}}，{entrypoint_requirement}。只输出完整 Python 代码。
         "numpy",
         "app.backtest.matrix",
         "app.strategy.builtin.factor_rank_research",
+        "app.strategy.market_data",   # 新增: 策略可读取指数/ETF 日K
         "datetime",
         "__future__",
     })

backend/app/strategy/market_data.py

@@ -0,0 +1,130 @@
+"""策略可访问的指数/ETF 日K读取模块 — 白名单放行的只读数据入口。
+
+供 Custom/AI 策略在 filter_history 内读取任意指数(及 ETF)的完整日K。
+策略通过白名单 import 本模块, 调用纯读函数; 禁止写操作或任意文件访问。
+
+设计要点:
+  - 模块自身是框架侧信任代码, 对策略的沙箱逃逸拦截(ai_generator._validate_safety)照旧生效。
+  - repo 线程安全懒加载(首次调用才构建); DataStore() 默认 settings.data_dir, 与 main.py 同源。
+  - 未知 symbol / 数据缺失 → 返回空 DataFrame(不抛), 与 repo 语义一致。
+"""
+from __future__ import annotations
+
+import logging
+import threading
+from datetime import date
+from typing import Any
+
+import polars as pl
+
+logger = logging.getLogger(__name__)
+
+# 完整历史默认区间下界(A股数据远晚于此, 仅作"全量"占位)。
+_FULL_START = date(1990, 1, 1)
+
+# ── repo 懒加载(线程安全) ─────────────────────────────
+_repo = None
+_lock = threading.Lock()
+
+
+def _get_repo():
+    global _repo
+    if _repo is None:
+        with _lock:
+            if _repo is None:
+                from app.tickflow.repository import DataStore, KlineRepository
+                _repo = KlineRepository(DataStore())
+    return _repo
+
+
+def _set_repo(repo: Any) -> None:
+    """测试注入: 用 fake repo 替换单例。"""
+    global _repo
+    with _lock:
+        _repo = repo
+
+
+def _reset_repo() -> None:
+    """测试清理: 重置单例, 下次调用重新懒加载。"""
+    global _repo
+    with _lock:
+        _repo = None
+
+
+# ── 参数规范化 ─────────────────────────────────────────
+def _norm_date(value, default: date) -> date:
+    if value is None:
+        return default
+    if isinstance(value, str):
+        return date.fromisoformat(value)
+    return value
+
+
+def _validate_symbol(symbol: Any) -> bool:
+    return isinstance(symbol, str) and bool(symbol.strip())
+
+
+# ── 公开只读 API ───────────────────────────────────────
+def get_index_daily(symbol, start=None, end=None, columns=None):
+    """读取指数日K(含技术指标)。未知 symbol / 无数据返回空 DataFrame。"""
+    if not _validate_symbol(symbol):
+        logger.warning("market_data: 非法指数 symbol %r", symbol)
+        return pl.DataFrame()
+    s = _norm_date(start, _FULL_START)
+    e = _norm_date(end, date.today())
+    try:
+        return _get_repo().get_index_daily(symbol, s, e, columns)
+    except Exception as exc:
+        logger.warning("market_data get_index_daily failed %s: %s", symbol, exc)
+        return pl.DataFrame()
+
+
+def get_etf_daily(symbol, start=None, end=None, columns=None):
+    """读取 ETF 日K(含技术指标)。同 get_index_daily 语义。"""
+    if not _validate_symbol(symbol):
+        logger.warning("market_data: 非法 ETF symbol %r", symbol)
+        return pl.DataFrame()
+    s = _norm_date(start, _FULL_START)
+    e = _norm_date(end, date.today())
+    try:
+        return _get_repo().get_etf_daily(symbol, s, e, columns)
+    except Exception as exc:
+        logger.warning("market_data get_etf_daily failed %s: %s", symbol, exc)
+        return pl.DataFrame()
+
+
+def get_daily(symbol, start=None, end=None, columns=None):
+    """按资产类型自动分派读取日K: 指数 → get_index_daily; ETF → get_etf_daily; 股票 → get_daily。"""
+    if not _validate_symbol(symbol):
+        logger.warning("market_data: 非法 symbol %r", symbol)
+        return pl.DataFrame()
+    s = _norm_date(start, _FULL_START)
+    e = _norm_date(end, date.today())
+    repo = _get_repo()
+    try:
+        asset_type = repo.resolve_asset_type(symbol)
+        if asset_type == "index":
+            return repo.get_index_daily(symbol, s, e, columns)
+        if asset_type == "etf":
+            return repo.get_etf_daily(symbol, s, e, columns)
+        return repo.get_daily(symbol, s, e, columns)
+    except Exception as exc:
+        logger.warning("market_data get_daily failed %s: %s", symbol, exc)
+        return pl.DataFrame()
+
+
+def list_index_symbols() -> list[dict]:
+    """列出已收录的指数符号(含名称)。无数据返回空列表。"""
+    try:
+        df = _get_repo().get_instruments_asset("index")
+    except Exception as exc:
+        logger.warning("market_data list_index_symbols failed: %s", exc)
+        return []
+    if df.is_empty() or "symbol" not in df.columns:
+        return []
+    name_col = "name" if "name" in df.columns else None
+    cols = ["symbol"] + ([name_col] if name_col else [])
+    return [
+        {"symbol": row["symbol"], "name": row.get("name")}
+        for row in df.select(cols).iter_rows(named=True)
+    ]

backend/tests/test_strategy_market_data.py

@@ -0,0 +1,101 @@
+"""策略指数K线访问模块 — 测试。"""
+import datetime
+
+import polars as pl
+import pytest
+
+from app.strategy import market_data
+from app.strategy.ai_generator import AIStrategyGenerator
+
+
+def test_whitelist_allows_market_data_import():
+    AIStrategyGenerator._validate_safety(
+        "from app.strategy.market_data import get_index_daily, get_daily"
+    )
+
+
+def test_whitelist_still_blocks_dangerous():
+    with pytest.raises(ValueError):
+        AIStrategyGenerator._validate_safety("import os")
+    with pytest.raises(ValueError):
+        AIStrategyGenerator._validate_safety("from os import path")
+    with pytest.raises(ValueError):
+        AIStrategyGenerator._validate_safety("getattr(obj, '__globals__')")
+
+
+class _FakeRepo:
+    """最小 fake: 只实现 market_data 用到的接口。"""
+    def __init__(self, index_df=None):
+        self.calls: list[tuple] = []
+        self._asset = {"000001.SH": "index", "510300.SH": "etf", "600000.SH": "stock"}
+        self._index_df = index_df if index_df is not None else pl.DataFrame(
+            {"date": ["2026-01-02"], "close": [3000.0], "macd_dif": [1.0], "macd_dea": [2.0]}
+        )
+        self._empty = pl.DataFrame()
+
+    def resolve_asset_type(self, symbol):
+        self.calls.append(("resolve", symbol))
+        return self._asset.get(symbol, "stock")
+
+    def get_index_daily(self, symbol, start=None, end=None, columns=None):
+        self.calls.append(("index", symbol, start, end, columns))
+        return self._index_df if symbol == "000001.SH" else self._empty
+
+    def get_etf_daily(self, symbol, start=None, end=None, columns=None):
+        self.calls.append(("etf", symbol, start, end, columns))
+        return self._empty
+
+    def get_daily(self, symbol, start=None, end=None, columns=None):
+        self.calls.append(("stock", symbol, start, end, columns))
+        return self._empty
+
+    def get_instruments_asset(self, asset_type):
+        return pl.DataFrame({"symbol": ["000001.SH"], "name": ["上证指数"]})
+
+
+@pytest.fixture()
+def fake_repo():
+    fake = _FakeRepo()
+    market_data._set_repo(fake)
+    yield fake
+    market_data._reset_repo()
+
+
+def test_get_index_daily_delegates_and_normalizes_dates(fake_repo):
+    df = market_data.get_index_daily(
+        "000001.SH", start="2026-01-01", end="2026-01-31", columns=["date", "close"]
+    )
+    assert df.height == 1 and df["close"][0] == 3000.0
+    _, sym, s, e, cols = fake_repo.calls[-1]
+    assert sym == "000001.SH"
+    assert s == datetime.date(2026, 1, 1)
+    assert e == datetime.date(2026, 1, 31)
+    assert cols == ["date", "close"]
+
+
+@pytest.mark.parametrize("symbol,expected_kind", [
+    ("000001.SH", "index"),
+    ("510300.SH", "etf"),
+    ("600000.SH", "stock"),
+])
+def test_get_daily_dispatch_by_asset_type(fake_repo, symbol, expected_kind):
+    market_data.get_daily(symbol)
+    last = fake_repo.calls[-1]
+    assert last[0] == expected_kind
+    assert last[1] == symbol
+
+
+def test_bad_symbol_returns_empty_without_calling_repo(fake_repo):
+    assert market_data.get_index_daily("").is_empty()
+    assert market_data.get_index_daily(None).is_empty()
+    assert market_data.get_etf_daily("").is_empty()
+    assert market_data.get_daily(None).is_empty()
+    assert fake_repo.calls == []
+
+
+def test_missing_symbol_returns_empty_no_raise(fake_repo):
+    assert market_data.get_index_daily("999999.SH").is_empty()
+
+
+def test_list_index_symbols(fake_repo):
+    assert market_data.list_index_symbols() == [{"symbol": "000001.SH", "name": "上证指数"}]