""" Python 策略脚本(on_init / on_bar + ctx.buy/sell/close_position)运行时。 与回测逻辑对齐,供 TradingExecutor 实盘逐根 K 线调用。 """ from __future__ import annotations from typing import Any, Callable, Dict, List, Optional, Tuple import numpy as np import pandas as pd from app.utils.logger import get_logger logger = get_logger(__name__) class ScriptBar(dict): def __getattr__(self, name: str) -> Any: try: return self[name] except KeyError as exc: raise AttributeError(name) from exc class ScriptPosition(dict): def __init__(self): super().__init__() self.clear_position() def __getattr__(self, name: str) -> Any: try: return self[name] except KeyError as exc: raise AttributeError(name) from exc def __bool__(self) -> bool: return bool(self.get("side")) and float(self.get("size") or 0) > 0 def __int__(self) -> int: return int(self.get("direction") or 0) def __float__(self) -> float: return float(self.get("direction") or 0) def __eq__(self, other: Any) -> bool: try: return int(self) == int(other) except Exception: return dict.__eq__(self, other) def __lt__(self, other: Any) -> bool: return int(self) < int(other) def __le__(self, other: Any) -> bool: return int(self) <= int(other) def __gt__(self, other: Any) -> bool: return int(self) > int(other) def __ge__(self, other: Any) -> bool: return int(self) >= int(other) def clear_position(self) -> None: self.clear() self.update( { "side": "", "size": 0.0, "entry_price": 0.0, "direction": 0, "amount": 0.0, } ) def open_position(self, side: str, entry_price: float, amount: float) -> None: direction = 1 if side == "long" else (-1 if side == "short" else 0) size = float(amount or 0.0) price = float(entry_price or 0.0) self.clear() self.update( { "side": side, "size": size, "entry_price": price, "direction": direction, "amount": size, } ) def add_position(self, entry_price: float, amount: float) -> None: extra = float(amount or 0.0) if extra <= 0: return current_size = float(self.get("size") or 0.0) current_price = float(self.get("entry_price") or 0.0) next_size = current_size + extra next_price = float(entry_price or current_price or 0.0) if current_size > 0 and current_price > 0 and next_size > 0: next_price = ((current_price * current_size) + (float(entry_price or current_price) * extra)) / next_size self["size"] = next_size self["amount"] = next_size self["entry_price"] = next_price class StrategyScriptContext: """与回测 ScriptBacktestContext 行为一致,供实盘按根推进。""" def __init__(self, bars_df: pd.DataFrame, initial_balance: float): self._bars_df = bars_df self._params: Dict[str, Any] = {} self._orders: List[Dict[str, Any]] = [] self._logs: List[str] = [] self.current_index = -1 self.position = ScriptPosition() self.balance = float(initial_balance) self.equity = float(initial_balance) def param(self, name: str, default: Any = None) -> Any: if name not in self._params: self._params[name] = default return self._params[name] def bars(self, n: int = 1): start = max(0, self.current_index - int(n) + 1) out = [] for _, row in self._bars_df.iloc[start : self.current_index + 1].iterrows(): out.append( ScriptBar( open=float(row.get("open") or 0), high=float(row.get("high") or 0), low=float(row.get("low") or 0), close=float(row.get("close") or 0), volume=float(row.get("volume") or 0), timestamp=row.get("time"), ) ) return out def log(self, message: Any): self._logs.append(str(message)) def buy(self, price: Any = None, amount: Any = None): self._orders.append({"action": "buy", "price": price, "amount": amount}) def sell(self, price: Any = None, amount: Any = None): self._orders.append({"action": "sell", "price": price, "amount": amount}) def close_position(self): self._orders.append({"action": "close"}) def compile_strategy_script_handlers(code: str) -> Tuple[Optional[Callable], Optional[Callable]]: """ 校验并编译策略脚本,返回 (on_init, on_bar)。 on_bar 不可缺省;on_init 可选。 """ if not code or not str(code).strip(): raise ValueError("Strategy script is empty") import builtins def safe_import(name, *args, **kwargs): allowed_modules = ["numpy", "pandas", "math", "json", "datetime", "time"] if name in allowed_modules or name.split(".")[0] in allowed_modules: return builtins.__import__(name, *args, **kwargs) raise ImportError(f"Import not allowed: {name}") safe_builtins = { k: getattr(builtins, k) for k in dir(builtins) if not k.startswith("_") and k not in ["eval", "exec", "compile", "open", "input", "help", "exit", "quit"] } safe_builtins["__import__"] = safe_import exec_env = { "__builtins__": safe_builtins, "np": np, "pd": pd, } from app.utils.safe_exec import safe_exec_code, validate_code_safety is_safe, error_msg = validate_code_safety(code) if not is_safe: raise ValueError(f"Code contains unsafe operations: {error_msg}") exec_result = safe_exec_code(code=code, exec_globals=exec_env, exec_locals=exec_env, timeout=60) if not exec_result["success"]: raise RuntimeError(f"Code execution failed: {exec_result.get('error')}") on_init = exec_env.get("on_init") on_bar = exec_env.get("on_bar") if not callable(on_bar): raise ValueError("Strategy script must define on_bar(ctx, bar)") if on_init is not None and not callable(on_init): on_init = None return (on_init if callable(on_init) else None), on_bar