staplebridge/data/dataset.py

from __future__ import annotations

import json
from pathlib import Path

from staplebridge.data.schemas import LeadExample


def save_leads(leads: list[LeadExample], out_path: str | Path) -> None:
    path = Path(out_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        for lead in leads:
            f.write(json.dumps(lead.to_dict()) + "\n")


def load_leads(path: str | Path) -> list[LeadExample]:
    leads: list[LeadExample] = []
    with Path(path).open("r", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                leads.append(LeadExample(**json.loads(line)))
    return leads