224 lines
8.2 KiB
Python
224 lines
8.2 KiB
Python
"""Enumerate vnpy_ctastrategy CTA strategies + their parameters.
|
||
|
||
Used by the backtest UI dropdown and dynamic parameter form. Falls back to a
|
||
static name list when vnpy_ctastrategy is not importable (e.g. local dev).
|
||
Task S1.3.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import ast
|
||
import glob
|
||
import importlib
|
||
import logging
|
||
import os
|
||
import pkgutil
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Fallback strategy names (when vnpy_ctastrategy import fails).
|
||
STRATEGY_NAMES: list[str] = ["DoubleMaStrategy", "BollChannelStrategy", "AtrRsiStrategy"]
|
||
|
||
|
||
def _load_strategy_classes() -> dict[str, type]:
|
||
"""Import all Strategy classes from vnpy_ctastrategy.strategies."""
|
||
classes: dict[str, type] = {}
|
||
try:
|
||
mod = importlib.import_module("vnpy_ctastrategy.strategies")
|
||
for _, name, _ in pkgutil.iter_modules(mod.__path__):
|
||
try:
|
||
m = importlib.import_module(f"vnpy_ctastrategy.strategies.{name}")
|
||
for attr in dir(m):
|
||
obj = getattr(m, attr)
|
||
if isinstance(obj, type) and attr.endswith("Strategy") and hasattr(obj, "parameters"):
|
||
classes[attr] = obj
|
||
except Exception as e:
|
||
logger.warning("导入策略模块 %s 失败: %s", name, e)
|
||
continue
|
||
except Exception as e:
|
||
logger.warning("加载 vnpy_ctastrategy 策略列表失败（降级为静态列表）: %s", e)
|
||
return classes
|
||
|
||
|
||
def list_strategies() -> list[dict]:
|
||
"""Return [{name, class_name}, ...] for the UI dropdown.
|
||
|
||
vnpy_ctastrategy 导入的 CTA 类 + 自研目录 AST 扫出的 CTA 类名（不 import，
|
||
避免重依赖副作用）。全空时降级 STRATEGY_NAMES。
|
||
"""
|
||
classes = _load_strategy_classes()
|
||
items = [{"name": n, "class_name": n} for n in sorted(classes)]
|
||
known = {i["name"] for i in items}
|
||
for f in _scan_self_owned():
|
||
if f["type"] == "cta" and f["class_name"] not in known:
|
||
items.append({"name": f["class_name"], "class_name": f["class_name"]})
|
||
known.add(f["class_name"])
|
||
if items:
|
||
return items
|
||
return [{"name": n, "class_name": n} for n in STRATEGY_NAMES]
|
||
|
||
|
||
def strategy_params(name: str) -> dict:
|
||
"""Return {parameters: [...], defaults: {...}} for a strategy's dynamic form."""
|
||
cls = _load_self_owned_module_classes().get(name) or _load_strategy_classes().get(name)
|
||
if cls is None:
|
||
return {"parameters": [], "defaults": {}}
|
||
params = list(getattr(cls, "parameters", []))
|
||
defaults = {p: getattr(cls, p, None) for p in params}
|
||
return {"parameters": params, "defaults": defaults}
|
||
|
||
|
||
def _load_self_owned_module_classes() -> dict[str, type]:
|
||
"""从自研 CTA 目录(sanguo_trader/strategy)动态 import 策略类。
|
||
|
||
让「复制进来的 vnpy 内置模板 + 在线编辑的自研策略」真正可被回测加载,
|
||
且优先于 pip 包同名类(编辑副本即刻生效)。import 失败(如本机 dev 无
|
||
vnpy_ctastrategy)静默降级到 pip 包。
|
||
"""
|
||
classes: dict[str, type] = {}
|
||
pkg = SELF_OWNED_DIRS[0].replace("/", ".")
|
||
for f in _scan_self_owned():
|
||
if f["type"] != "cta" or f["dir"] != SELF_OWNED_DIRS[0] + "/":
|
||
continue
|
||
stem = f["name"][:-3]
|
||
try:
|
||
mod = importlib.import_module(f"{pkg}.{stem}")
|
||
cls = getattr(mod, f["class_name"], None)
|
||
if isinstance(cls, type):
|
||
classes[f["class_name"]] = cls
|
||
except Exception as e:
|
||
logger.debug("自研策略 %s 导入失败(降级 pip 包): %s", f["name"], e)
|
||
return classes
|
||
|
||
|
||
def get_strategy_class(name: str) -> type | None:
|
||
"""Return the strategy class by name (None if unavailable).
|
||
|
||
自研目录优先(编辑生效),pip 包兜底。
|
||
"""
|
||
cls = _load_self_owned_module_classes().get(name)
|
||
if cls is not None:
|
||
return cls
|
||
return _load_strategy_classes().get(name)
|
||
|
||
|
||
# ===== 自研策略目录扫描 + 文件读写（在线编辑，spec §12）=====
|
||
|
||
# 自研策略目录（相对仓库根 / 容器 /app）。sanguo_live/strategies 不存在，勿加。
|
||
SELF_OWNED_DIRS: list[str] = [
|
||
"sanguo_trader/strategy",
|
||
"sanguo_portfolio/strategies",
|
||
]
|
||
|
||
|
||
def _repo_root() -> str:
|
||
"""strategy_registry.py 所在包的上一级 = 仓库根（/app 或开发目录）。"""
|
||
return os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
|
||
def _classify_strategy_type(tree: ast.Module) -> str:
|
||
"""AST 判定 portfolio vs cta：
|
||
- import 了 bullet_trade 风格 StrategyTemplate → cta
|
||
- from vnpy_ctastrategy import ...（vnpy 内置/自研 CTA 文件风格）→ cta
|
||
否则 portfolio。"""
|
||
for node in ast.walk(tree):
|
||
if isinstance(node, ast.ImportFrom):
|
||
mod = node.module or ""
|
||
if mod.endswith("StrategyTemplate") or mod.endswith("vnpy_ctastrategy"):
|
||
return "cta"
|
||
if isinstance(node, ast.Import):
|
||
for alias in node.names:
|
||
if "StrategyTemplate" in alias.name or "vnpy_ctastrategy" in alias.name:
|
||
return "cta"
|
||
return "portfolio"
|
||
|
||
|
||
def _extract_class_names(tree: ast.Module) -> list[str]:
|
||
"""AST 提取所有 `class XxxStrategy` 类名。"""
|
||
return [
|
||
n.name for n in ast.walk(tree)
|
||
if isinstance(n, ast.ClassDef) and n.name.endswith("Strategy")
|
||
]
|
||
|
||
|
||
def _scan_self_owned() -> list[dict]:
|
||
"""AST 扫描自研目录（不 import，避免重依赖副作用）。"""
|
||
root = _repo_root()
|
||
out: list[dict] = []
|
||
for rel in SELF_OWNED_DIRS:
|
||
abs_dir = os.path.join(root, rel)
|
||
for py in sorted(glob.glob(os.path.join(abs_dir, "*.py"))):
|
||
name = os.path.basename(py)
|
||
if name.startswith("_"):
|
||
continue
|
||
try:
|
||
src = open(py, encoding="utf-8").read()
|
||
tree = ast.parse(src, filename=name)
|
||
except (SyntaxError, OSError):
|
||
continue
|
||
classes = _extract_class_names(tree)
|
||
if not classes:
|
||
continue
|
||
out.append({
|
||
"name": name,
|
||
"dir": rel + "/",
|
||
"class_name": classes[0],
|
||
"type": _classify_strategy_type(tree),
|
||
"lines": len(src.splitlines()),
|
||
"modified": "",
|
||
})
|
||
return out
|
||
|
||
|
||
def list_strategy_files() -> dict:
|
||
"""GET /strategy/files → {files: [...]}（不含 code，前端列表用）。"""
|
||
return {"files": _scan_self_owned()}
|
||
|
||
|
||
def read_strategy_file(name: str) -> dict:
|
||
"""按文件名读单个策略文件全量（含 code）。找不到 → ValueError。"""
|
||
root = _repo_root()
|
||
for rel in SELF_OWNED_DIRS:
|
||
path = os.path.join(root, rel, name)
|
||
if os.path.exists(path):
|
||
src = open(path, encoding="utf-8").read()
|
||
tree = ast.parse(src, filename=name)
|
||
classes = _extract_class_names(tree)
|
||
return {
|
||
"name": name,
|
||
"dir": rel + "/",
|
||
"class_name": classes[0] if classes else "",
|
||
"type": _classify_strategy_type(tree),
|
||
"lines": len(src.splitlines()),
|
||
"modified": "",
|
||
"code": src,
|
||
}
|
||
raise ValueError(f"策略文件不存在: {name}")
|
||
|
||
|
||
def write_strategy_file(name: str, code: str) -> None:
|
||
"""写策略文件。py_compile 语法校验失败 → ValueError（拒绝坏代码落盘）。"""
|
||
import py_compile
|
||
import tempfile
|
||
with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False, encoding="utf-8") as tf:
|
||
tf.write(code)
|
||
tmp = tf.name
|
||
try:
|
||
py_compile.compile(tmp, doraise=True)
|
||
except py_compile.PyCompileError as e:
|
||
raise ValueError(f"语法错误，未保存: {e}") from e
|
||
finally:
|
||
os.unlink(tmp)
|
||
|
||
root = _repo_root()
|
||
for rel in SELF_OWNED_DIRS:
|
||
path = os.path.join(root, rel, name)
|
||
if os.path.exists(path):
|
||
with open(path, "w", encoding="utf-8") as f:
|
||
f.write(code)
|
||
return
|
||
# 新文件落 cta 目录（MVP 统一落 sanguo_trader/strategy）
|
||
target = os.path.join(root, SELF_OWNED_DIRS[0], name)
|
||
os.makedirs(os.path.dirname(target), exist_ok=True)
|
||
with open(target, "w", encoding="utf-8") as f:
|
||
f.write(code)
|