132 lines
4.0 KiB
Python
132 lines
4.0 KiB
Python
"""
|
|
Excel / CSV / JSON export.
|
|
|
|
Excel uses openpyxl via pandas — multi-sheet workbooks readable in
|
|
Excel 2019+. JSON round-trips the full input state (sites, takeouts,
|
|
feature scopes) so a scenario can be saved and reloaded.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pandas as pd
|
|
|
|
from .inputs import CostTakeout, FeatureScope, SiteInput
|
|
from .meters import Confidence, TokenMeter
|
|
from .rollout import RolloutPlan
|
|
|
|
|
|
def meters_dataframe(meters: dict[str, TokenMeter]) -> pd.DataFrame:
|
|
"""Meter catalogue as a display/export-ready DataFrame."""
|
|
return pd.DataFrame(
|
|
[
|
|
{
|
|
"feature": m.feature,
|
|
"meter_type": m.meter_type.value,
|
|
"units_per_token": m.units_per_token or None,
|
|
"tokens_per_unit": m.tokens_per_unit,
|
|
"confidence": f"{m.confidence.icon} {m.confidence.value}",
|
|
"notes": m.notes,
|
|
"source": m.source_url or "",
|
|
}
|
|
for m in meters.values()
|
|
]
|
|
)
|
|
|
|
|
|
def sites_dataframe(sites: list[SiteInput]) -> pd.DataFrame:
|
|
rows = []
|
|
for s in sites:
|
|
d = dataclasses.asdict(s)
|
|
d["languages"] = ", ".join(d["languages"])
|
|
rows.append(d)
|
|
return pd.DataFrame(rows)
|
|
|
|
|
|
def export_excel(
|
|
sheets: dict[str, pd.DataFrame],
|
|
path: str | Path,
|
|
) -> Path:
|
|
"""Write a multi-sheet Excel workbook. Sheet names are truncated to
|
|
Excel's 31-character limit."""
|
|
path = Path(path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with pd.ExcelWriter(path, engine="openpyxl") as writer:
|
|
for name, df in sheets.items():
|
|
df.to_excel(writer, sheet_name=name[:31], index=False)
|
|
return path
|
|
|
|
|
|
def export_csv(df: pd.DataFrame, path: str | Path) -> Path:
|
|
path = Path(path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
df.to_csv(path, index=False)
|
|
return path
|
|
|
|
|
|
# ── JSON scenario save / load ────────────────────────────────────────
|
|
|
|
def scenario_state_to_json(
|
|
sites: list[SiteInput],
|
|
takeouts: list[CostTakeout],
|
|
feature_scopes: list[FeatureScope],
|
|
path: str | Path | None = None,
|
|
rollout: RolloutPlan | None = None,
|
|
) -> str:
|
|
"""Serialize the full input state; optionally write to ``path``."""
|
|
state = {
|
|
"sites": [dataclasses.asdict(s) for s in sites],
|
|
"takeouts": [
|
|
{**dataclasses.asdict(t), "confidence": t.confidence.value}
|
|
for t in takeouts
|
|
],
|
|
"feature_scopes": [
|
|
{
|
|
**dataclasses.asdict(f),
|
|
"adoption_curve": {str(k): v for k, v in f.adoption_curve.items()},
|
|
}
|
|
for f in feature_scopes
|
|
],
|
|
}
|
|
if rollout is not None:
|
|
state["rollout"] = dataclasses.asdict(rollout)
|
|
text = json.dumps(state, indent=2)
|
|
if path is not None:
|
|
Path(path).write_text(text)
|
|
return text
|
|
|
|
|
|
def scenario_state_from_json(
|
|
source: str | Path,
|
|
) -> tuple[list[SiteInput], list[CostTakeout], list[FeatureScope], RolloutPlan | None]:
|
|
"""Inverse of :func:`scenario_state_to_json`. ``source`` is a JSON
|
|
string or a file path. The fourth element is None for legacy files
|
|
saved without a rollout plan."""
|
|
raw = (
|
|
Path(source).read_text()
|
|
if isinstance(source, Path) or (isinstance(source, str) and source.strip().endswith(".json"))
|
|
else str(source)
|
|
)
|
|
state = json.loads(raw)
|
|
sites = [SiteInput(**s) for s in state["sites"]]
|
|
takeouts = [
|
|
CostTakeout(**{**t, "confidence": Confidence(t["confidence"])})
|
|
for t in state["takeouts"]
|
|
]
|
|
scopes = [
|
|
FeatureScope(
|
|
**{
|
|
**f,
|
|
"adoption_curve": {int(k): v for k, v in f["adoption_curve"].items()},
|
|
}
|
|
)
|
|
for f in state["feature_scopes"]
|
|
]
|
|
rollout = (
|
|
RolloutPlan(**state["rollout"]) if "rollout" in state else None
|
|
)
|
|
return sites, takeouts, scopes, rollout
|