MoFin/venv/lib/python3.12/site-packages/alphasift/store.py

# -*- coding: utf-8 -*-
"""Persistence helpers for screen runs and evaluations."""

from __future__ import annotations

import json
from dataclasses import asdict, fields
from pathlib import Path

from alphasift.models import EvaluationResult, Pick, PickEvaluation, ScreenResult

_RUN_METADATA_SUFFIX = ".meta"


def save_screen_result(
    result: ScreenResult,
    *,
    data_dir: Path,
    path: str | Path | None = None,
    jsonl: bool = False,
) -> Path:
    """Persist a screen result and return the written path."""
    output_path = Path(path) if path is not None else data_dir / "runs" / f"{result.run_id}.json"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    result.saved_path = str(output_path)
    if jsonl:
        output_path.write_text("\n".join(screen_result_to_jsonl(result)) + "\n", encoding="utf-8")
    else:
        output_path.write_text(
            json.dumps(asdict(result), ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
    _write_screen_result_metadata(result, output_path)
    return output_path


def load_screen_result(run_ref: str | Path, *, data_dir: Path) -> ScreenResult:
    """Load a saved screen result by run_id or path."""
    path = resolve_run_path(run_ref, data_dir=data_dir)
    data = json.loads(path.read_text(encoding="utf-8"))
    pick_items = data.get("picks", [])
    pick_fields = {field.name for field in fields(Pick)}
    result_fields = {field.name for field in fields(ScreenResult)}
    data["picks"] = [
        Pick(**{key: value for key, value in item.items() if key in pick_fields})
        for item in pick_items
        if isinstance(item, dict)
    ]
    filtered = {key: value for key, value in data.items() if key in result_fields}
    loaded = ScreenResult(**filtered)
    loaded.saved_path = str(path)
    return loaded


def save_evaluation_result(
    result: EvaluationResult,
    *,
    data_dir: Path,
    path: str | Path | None = None,
    jsonl: bool = False,
) -> Path:
    output_path = Path(path) if path is not None else data_dir / "evaluations" / f"{result.run_id}.json"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    result.saved_path = str(output_path)
    if jsonl:
        output_path.write_text("\n".join(evaluation_result_to_jsonl(result)) + "\n", encoding="utf-8")
    else:
        output_path.write_text(
            json.dumps(asdict(result), ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
    return output_path


def list_saved_runs(
    *,
    data_dir: Path,
    limit: int = 20,
    strategy: str | None = None,
) -> list[dict[str, object]]:
    runs_dir = data_dir / "runs"
    if not runs_dir.is_dir():
        return []
    limit = int(limit)
    if limit <= 0:
        return []
    items = []
    for path in sorted(runs_dir.glob("*.json"), key=lambda p: p.stat().st_mtime, reverse=True):
        metadata = _read_run_metadata(path)
        if metadata is None:
            continue
        if strategy and metadata.get("strategy") != strategy:
            continue
        items.append(metadata)
        if len(items) >= limit:
            break
    return items


def _write_screen_result_metadata(result: ScreenResult, output_path: Path) -> None:
    metadata_path = _run_metadata_path(output_path)
    metadata = {
        "run_id": result.run_id or output_path.stem,
        "strategy": result.strategy,
        "market": result.market,
        "created_at": result.created_at,
        "picks": len(result.picks),
        "path": str(output_path),
    }
    metadata_path.write_text(
        json.dumps(metadata, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )


def _run_metadata_path(run_path: Path) -> Path:
    return run_path.with_suffix(run_path.suffix + _RUN_METADATA_SUFFIX)


def _read_run_metadata(path: Path) -> dict[str, object] | None:
    metadata_path = _run_metadata_path(path)
    if metadata_path.is_file():
        try:
            metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
            if isinstance(metadata, dict):
                return _normalize_run_metadata(metadata, path)
        except Exception:
            pass
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return None
    if not isinstance(data, dict):
        return None
    return _normalize_run_metadata(data, path)


def _normalize_run_metadata(data: dict, path: Path) -> dict[str, object]:
    picks = data.get("picks", [])
    if isinstance(picks, list):
        pick_count = len(picks)
    else:
        try:
            pick_count = int(picks)
        except (TypeError, ValueError):
            pick_count = 0
    return {
        "run_id": data.get("run_id", path.stem),
        "strategy": data.get("strategy", ""),
        "market": data.get("market", ""),
        "created_at": data.get("created_at", ""),
        "picks": pick_count,
        "path": str(path),
    }


def resolve_run_path(run_ref: str | Path, *, data_dir: Path) -> Path:
    path = Path(run_ref)
    if path.is_file():
        return path
    candidate = data_dir / "runs" / f"{run_ref}.json"
    if candidate.is_file():
        return candidate
    raise FileNotFoundError(f"Saved run not found: {run_ref}")


def screen_result_to_jsonl(result: ScreenResult) -> list[str]:
    data = asdict(result)
    picks = data.pop("picks", [])
    lines = [json.dumps({"type": "run", **data}, ensure_ascii=False)]
    for pick in picks:
        lines.append(json.dumps({"type": "pick", "run_id": result.run_id, **pick}, ensure_ascii=False))
    return lines


def evaluation_result_to_jsonl(result: EvaluationResult) -> list[str]:
    data = asdict(result)
    picks = data.pop("picks", [])
    lines = [json.dumps({"type": "evaluation", **data}, ensure_ascii=False)]
    for pick in picks:
        lines.append(json.dumps({"type": "pick_evaluation", "run_id": result.run_id, **pick}, ensure_ascii=False))
    return lines


def evaluation_from_dict(data: dict) -> EvaluationResult:
    pick_fields = {field.name for field in fields(PickEvaluation)}
    result_fields = {field.name for field in fields(EvaluationResult)}
    data = dict(data)
    data["picks"] = [
        PickEvaluation(**{key: value for key, value in item.items() if key in pick_fields})
        for item in data.get("picks", [])
        if isinstance(item, dict)
    ]
    return EvaluationResult(**{key: value for key, value in data.items() if key in result_fields})