56 lines
2.0 KiB
Python
56 lines
2.0 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
class OpenLaneResultWriter:
|
|
def __init__(self, jsonl_path: str | Path, aggregated_path: str | Path) -> None:
|
|
self.jsonl_path = Path(jsonl_path)
|
|
self.aggregated_path = Path(aggregated_path)
|
|
|
|
def ensure_parent_dirs(self) -> None:
|
|
self.jsonl_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self.aggregated_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
def append_page(self, page: int, records: list[dict[str, Any]]) -> int:
|
|
self.ensure_parent_dirs()
|
|
written = 0
|
|
with self.jsonl_path.open("a", encoding="utf-8") as fh:
|
|
for record in records:
|
|
payload = {
|
|
"page": page,
|
|
"record": record,
|
|
}
|
|
fh.write(json.dumps(payload, ensure_ascii=False) + "\n")
|
|
written += 1
|
|
return written
|
|
|
|
def finalize(self, *, max_pages: int, completed_pages: list[int], failed_pages: list[int]) -> dict[str, Any]:
|
|
self.ensure_parent_dirs()
|
|
records: list[dict[str, Any]] = []
|
|
if self.jsonl_path.exists():
|
|
with self.jsonl_path.open("r", encoding="utf-8") as fh:
|
|
for line in fh:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
item = json.loads(line)
|
|
records.append(item)
|
|
|
|
summary = {
|
|
"max_pages": max_pages,
|
|
"completed_pages": sorted(completed_pages),
|
|
"failed_pages": sorted(failed_pages),
|
|
"total_pages_completed": len(set(completed_pages)),
|
|
"total_pages_failed": len(set(failed_pages)),
|
|
"total_records": len(records),
|
|
"items": records,
|
|
}
|
|
self.aggregated_path.write_text(
|
|
json.dumps(summary, ensure_ascii=False, indent=2),
|
|
encoding="utf-8",
|
|
)
|
|
return summary
|