feat: load backtest-results from either file or zip file

This commit is contained in:
Matthias
2024-12-16 19:42:32 +01:00
parent b1fc271a7e
commit d2d6c6fa81
2 changed files with 21 additions and 4 deletions

View File

@@ -3,8 +3,10 @@ Helpers when analyzing backtest data
"""
import logging
import zipfile
from copy import copy
from datetime import datetime, timezone
from io import StringIO
from pathlib import Path
from typing import Any, Literal
@@ -165,8 +167,23 @@ def load_backtest_stats(filename: Path | str) -> BacktestResultType:
if not filename.is_file():
raise ValueError(f"File {filename} does not exist.")
logger.info(f"Loading backtest result from {filename}")
with filename.open() as file:
data = json_load(file)
if filename.suffix == ".zip":
try:
with zipfile.ZipFile(filename) as zipf:
json_filename = filename.with_suffix(".json")
try:
with zipf.open(json_filename.name) as json_file:
# Need to convert to StringIO since json_load expects a text stream
data = json_load(StringIO(json_file.read().decode("utf-8")))
except KeyError:
# File not found in zip
raise ValueError(f"Could not find {json_filename.name} in {filename}")
except zipfile.BadZipFile:
raise ValueError(f"Bad zip file: {filename}")
else:
with filename.open() as file:
data = json_load(file)
# Legacy list format does not contain metadata.
if isinstance(data, dict):

View File

@@ -72,7 +72,7 @@ def store_backtest_results(
# Store latest backtest info separately
latest_filename = Path.joinpath(zip_filename.parent, LAST_BT_RESULT_FN)
file_dump_json(latest_filename, {"latest_backtest": str(zip_filename.name)})
file_dump_json(latest_filename, {"latest_backtest": str(zip_filename.name)}, log=False)
# Create zip file and add the files
with zipfile.ZipFile(zip_filename, "w", zipfile.ZIP_DEFLATED) as zipf:
@@ -83,7 +83,7 @@ def store_backtest_results(
}
stats_buf = StringIO()
dump_json_to_file(stats_buf, stats_copy)
zipf.writestr(json_filename, stats_buf.getvalue())
zipf.writestr(json_filename.name, stats_buf.getvalue())
# Add market change data if present
if market_change_data is not None: