Files

65 lines
2.0 KiB
Python

"""
策略注册表 — 自动发现本目录下所有策略模块。
新增策略只需:
1. 在本目录创建 xxx.py, 定义一个继承 strategies.base.Strategy 的类
2. 该模块顶层需有 `STRATEGY_CLASS = XxxStrategy` (大写常量名固定)
3. 重启后自动出现在 `python main.py list` 中
无需手动修改本文件。
"""
from __future__ import annotations
import importlib
import pkgutil
from typing import Dict, Type
from .base import Strategy, SignalResult
__all__ = ["Strategy", "SignalResult", "get_strategies", "get_strategy"]
# 策略模块约定: 每个模块需暴露 STRATEGY_CLASS 常量
_STRATEGY_ATTR = "STRATEGY_CLASS"
def _discover() -> Dict[str, Type[Strategy]]:
"""扫描本包下所有模块,收集 STRATEGY_CLASS 常量"""
registry: Dict[str, Type[Strategy]] = {}
for _finder, mod_name, _is_pkg in pkgutil.iter_modules(__path__):
if mod_name.startswith("_") or mod_name == "base":
continue
try:
mod = importlib.import_module(f"{__name__}.{mod_name}")
except Exception as e:
print(f" ⚠️ 加载策略模块 {mod_name} 失败: {e}")
continue
cls = getattr(mod, _STRATEGY_ATTR, None)
if cls is None:
continue
if not isinstance(cls, type) or not issubclass(cls, Strategy):
print(f" ⚠️ {mod_name}.STRATEGY_CLASS 不是 Strategy 子类, 跳过")
continue
registry[cls.name] = cls
return registry
_CACHE: Dict[str, Type[Strategy]] | None = None
def get_strategies() -> Dict[str, Type[Strategy]]:
"""返回 {策略名: 策略类} 字典,首次调用时懒加载"""
global _CACHE
if _CACHE is None:
_CACHE = _discover()
return _CACHE
def get_strategy(name: str) -> Strategy:
"""按名称实例化策略"""
strategies = get_strategies()
if name not in strategies:
available = ", ".join(sorted(strategies.keys())) or "(无可用策略)"
raise KeyError(f"未知策略 '{name}'。可用: {available}")
return strategies[name]()