diff --git a/sanguo_backtest/__init__.py b/sanguo_backtest/__init__.py new file mode 100644 index 0000000..03ffe8c --- /dev/null +++ b/sanguo_backtest/__init__.py @@ -0,0 +1 @@ +# Sanguo Backtest Module diff --git a/sanguo_backtest/result_store.py b/sanguo_backtest/result_store.py new file mode 100644 index 0000000..033e9e0 --- /dev/null +++ b/sanguo_backtest/result_store.py @@ -0,0 +1,178 @@ +"""Backtest result storage using SQLite + parquet files.""" +import sqlite3 +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Optional +import pandas as pd + + +@dataclass +class BacktestResult: + """Backtest result data structure.""" + task_id: str + type: str + status: str + strategy: str + symbol: str + params: dict + start: str + end: str + statistics: dict + equity_curve: Optional[pd.DataFrame] = None + trades: Optional[pd.DataFrame] = None + error_msg: Optional[str] = None + + +# SQLite schema for backtest stats +_SCHEMA = """CREATE TABLE IF NOT EXISTS backtest_stats ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + task_id TEXT, + type TEXT, + status TEXT, + strategy TEXT, + symbol TEXT, + params TEXT, + start TEXT, + end TEXT, + statistics TEXT, + equity_path TEXT, + trades_path TEXT, + error_msg TEXT, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP +);""" + + +def _connect(db_path: str) -> sqlite3.Connection: + """Create database connection and initialize schema.""" + conn = sqlite3.connect(db_path) + conn.executescript(_SCHEMA) + return conn + + +def save_result(result: BacktestResult, db_path: str, file_dir: Optional[str] = None) -> int: + """ + Save backtest result to database and optionally save DataFrame fields to parquet files. + + Args: + result: BacktestResult object to save + db_path: Path to SQLite database file + file_dir: Optional directory to save parquet files (equity_curve and trades) + + Returns: + int: The ID of the inserted record + """ + conn = _connect(db_path) + try: + equity_path = None + trades_path = None + + # Save DataFrames to parquet if file_dir is provided + if file_dir: + fdir = Path(file_dir) + fdir.mkdir(parents=True, exist_ok=True) + + if result.equity_curve is not None and not result.equity_curve.empty: + equity_path = str(fdir / f"{result.task_id}_equity.parquet") + result.equity_curve.to_parquet(equity_path) + + if result.trades is not None and not result.trades.empty: + trades_path = str(fdir / f"{result.task_id}_trades.parquet") + result.trades.to_parquet(trades_path) + + # Insert record into database + cur = conn.execute( + """INSERT INTO backtest_stats + (task_id, type, status, strategy, symbol, params, start, end, + statistics, equity_path, trades_path, error_msg) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + ( + result.task_id, + result.type, + result.status, + result.strategy, + result.symbol, + json.dumps(result.params), + result.start, + result.end, + json.dumps(result.statistics), + equity_path, + trades_path, + result.error_msg + ) + ) + conn.commit() + return cur.lastrowid + finally: + conn.close() + + +def load_result(rid: int, db_path: str) -> BacktestResult: + """ + Load backtest result by ID from database. + + Args: + rid: Result ID to load + db_path: Path to SQLite database file + + Returns: + BacktestResult: Loaded result object + + Raises: + KeyError: If result ID not found + """ + conn = _connect(db_path) + try: + row = conn.execute("SELECT * FROM backtest_stats WHERE id=?", (rid,)).fetchone() + if not row: + raise KeyError(f"result {rid} not found") + + # Get column names from table description + cols = [d[0] for d in conn.execute("SELECT * FROM backtest_stats LIMIT 0").description] + d = dict(zip(cols, row)) + + # Load parquet files if paths exist + equity = pd.read_parquet(d["equity_path"]) if d.get("equity_path") else None + trades = pd.read_parquet(d["trades_path"]) if d.get("trades_path") else None + + return BacktestResult( + task_id=d["task_id"], + type=d["type"], + status=d["status"], + strategy=d["strategy"], + symbol=d["symbol"], + params=json.loads(d["params"]), + start=d["start"], + end=d["end"], + statistics=json.loads(d["statistics"]) if d["statistics"] else {}, + equity_curve=equity, + trades=trades, + error_msg=d.get("error_msg") + ) + finally: + conn.close() + + +def list_results(type_filter: Optional[str] = None, db_path: str = "") -> list[BacktestResult]: + """ + List all backtest results, optionally filtered by type. + + Args: + type_filter: Optional filter for result type (e.g., 'cta', 'factor') + db_path: Path to SQLite database file + + Returns: + list[BacktestResult]: List of loaded result objects + """ + conn = _connect(db_path) + try: + query = "SELECT id FROM backtest_stats" + args = () + + if type_filter: + query = query + " WHERE type=?" + args = (type_filter,) + + return [load_result(row[0], db_path) for row in conn.execute(query, args).fetchall()] + finally: + conn.close() diff --git a/tests/backtest/__init__.py b/tests/backtest/__init__.py new file mode 100644 index 0000000..17dbd4f --- /dev/null +++ b/tests/backtest/__init__.py @@ -0,0 +1 @@ +# Backtest tests package diff --git a/tests/backtest/conftest.py b/tests/backtest/conftest.py new file mode 100644 index 0000000..61d509e --- /dev/null +++ b/tests/backtest/conftest.py @@ -0,0 +1,20 @@ +"""Pytest fixtures for backtest tests.""" +import pytest +import tempfile +import shutil +from pathlib import Path + + +@pytest.fixture +def temp_db_path(tmp_path): + """Create a temporary database path.""" + db_path = tmp_path / "test_backtest.db" + yield str(db_path) + + +@pytest.fixture +def temp_file_dir(tmp_path): + """Create a temporary file directory for parquet files.""" + file_dir = tmp_path / "parquet_files" + file_dir.mkdir(parents=True, exist_ok=True) + yield str(file_dir) diff --git a/tests/backtest/test_result_store.py b/tests/backtest/test_result_store.py new file mode 100644 index 0000000..e22c54a --- /dev/null +++ b/tests/backtest/test_result_store.py @@ -0,0 +1,127 @@ +"""Tests for sanguo_backtest.result_store module.""" +import pytest +import pandas as pd +from datetime import datetime +from sanguo_backtest.result_store import ( + BacktestResult, + save_result, + load_result, + list_results +) + + +def test_save_and_load_result(temp_db_path): + """Test saving and loading a backtest result without parquet files.""" + # Create a result without equity_curve and trades (file_dir=None by default) + result = BacktestResult( + task_id="test_task_001", + type="cta", + status="completed", + strategy="DualThrustStrategy", + symbol="IF2401", + params={"window": 20, "entry_threshold": 0.5}, + start="2024-01-01", + end="2024-03-31", + statistics={ + "total_return": 0.15, + "sharpe_ratio": 1.2, + "max_drawdown": -0.08, + "win_rate": 0.55 + } + ) + + # Save the result + result_id = save_result(result, temp_db_path) + assert result_id > 0 + assert isinstance(result_id, int) + + # Load the result + loaded_result = load_result(result_id, temp_db_path) + + # Verify all fields match + assert loaded_result.task_id == result.task_id + assert loaded_result.type == result.type + assert loaded_result.status == result.status + assert loaded_result.strategy == result.strategy + assert loaded_result.symbol == result.symbol + assert loaded_result.params == result.params + assert loaded_result.start == result.start + assert loaded_result.end == result.end + assert loaded_result.statistics == result.statistics + assert loaded_result.equity_curve is None + assert loaded_result.trades is None + assert loaded_result.error_msg is None + + +def test_list_results_filter(temp_db_path): + """Test listing results with type filter.""" + # Create multiple results of different types + cta_result = BacktestResult( + task_id="cta_task_001", + type="cta", + status="completed", + strategy="DualThrustStrategy", + symbol="IF2401", + params={"window": 20}, + start="2024-01-01", + end="2024-03-31", + statistics={"total_return": 0.15} + ) + + factor_result = BacktestResult( + task_id="factor_task_001", + type="factor", + status="completed", + strategy="Alpha158Strategy", + symbol="000001.SZ", + params={"factors": ["alpha001", "alpha002"]}, + start="2024-01-01", + end="2024-03-31", + statistics={"total_return": 0.20} + ) + + # Save both results + cta_id = save_result(cta_result, temp_db_path) + factor_id = save_result(factor_result, temp_db_path) + + # List all results + all_results = list_results(db_path=temp_db_path) + assert len(all_results) == 2 + + # List only CTA results + cta_results = list_results(type_filter="cta", db_path=temp_db_path) + assert len(cta_results) == 1 + assert cta_results[0].type == "cta" + assert cta_results[0].task_id == "cta_task_001" + + # List only factor results + factor_results = list_results(type_filter="factor", db_path=temp_db_path) + assert len(factor_results) == 1 + assert factor_results[0].type == "factor" + assert factor_results[0].task_id == "factor_task_001" + + +def test_failed_result_stores_error_msg(temp_db_path): + """Test that failed results store error messages.""" + # Create a failed result with error message + failed_result = BacktestResult( + task_id="failed_task_001", + type="cta", + status="failed", + strategy="DualThrustStrategy", + symbol="IF2401", + params={"window": 20}, + start="2024-01-01", + end="2024-03-31", + statistics={}, # Empty statistics for failed result + error_msg="Data loading failed: insufficient historical data" + ) + + # Save the failed result + result_id = save_result(failed_result, temp_db_path) + + # Load and verify error message is preserved + loaded_result = load_result(result_id, temp_db_path) + assert loaded_result.status == "failed" + assert loaded_result.error_msg == "Data loading failed: insufficient historical data" + assert loaded_result.statistics == {}