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

Custom/AI 策略此前被沙箱隔离, filter_history(df, params) 只能拿到个股历史窗口,
无法读取指数K线, 也就做不了「大盘指数 MACD 死叉」这类市场级过滤。

本改动新增框架侧受信模块 backend/app/strategy/market_data.py, 暴露只读纯函数:
  - get_index_daily / get_etf_daily: 读取指数/ETF 日K(含技术指标, 支持列下推)
  - get_daily: 按 repo.resolve_asset_type 自动分派(指数/ETF/股票)
  - list_index_symbols: 枚举已收录指数
模块线程安全懒加载 repo (DataStore() 默认 settings.data_dir, 与 main.py 同源),
未知 symbol/缺数据返回空 DataFrame 不抛; 不向策略暴露文件访问或写能力。

ai_generator._ALLOWED_IMPORT_MODULES 放行 "app.strategy.market_data",
使 AI 生成与磁盘 Custom 策略均可 import 该模块; 其余模块/危险调用照旧拦截。

附 9 项框架级测试: 白名单放行/拦截、注入 fake repo 后委托与日期规范化、
按资产类型分派、坏 symbol/缺数据返回空、list_index_symbols。对既有策略零影响。

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
shuo-huge
2026-08-26 09:36:06 +08:00
co-authored by Claude Fable 5
parent 196af2fa06
commit 2f7a4b0979
3 changed files with 232 additions and 0 deletions
+1
View File
@@ -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__",
})
+130
View File
@@ -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)
]
+101
View File
@@ -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": "上证指数"}]