62 lines
2.3 KiB
Python
62 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Merge audit JSONL manifests into one unique source-ordered file."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
|
|
|
|
def main() -> None:
|
|
"""Merge the last result for each sample and write it in source order."""
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("source", type=Path, help="Source sample JSONL defining output order")
|
|
parser.add_argument("output", type=Path)
|
|
parser.add_argument("manifests", type=Path, nargs="+")
|
|
parser.add_argument("--require-count", type=int)
|
|
args = parser.parse_args()
|
|
|
|
audits: dict[str, dict] = {}
|
|
for path in args.manifests:
|
|
if not path.exists():
|
|
continue
|
|
with path.open(encoding="utf-8") as handle:
|
|
for line in handle:
|
|
if not line.strip():
|
|
continue
|
|
value = json.loads(line)
|
|
sample_id = value.get("sample_id")
|
|
if not isinstance(sample_id, str) or not sample_id:
|
|
raise ValueError(f"Audit in {path} has no sample_id")
|
|
audits[sample_id] = value
|
|
|
|
ordered: list[dict] = []
|
|
with args.source.open(encoding="utf-8") as handle:
|
|
for line in handle:
|
|
if not line.strip():
|
|
continue
|
|
source_record = json.loads(line)
|
|
sample_id = source_record.get("trajectory_id") or source_record.get("sample_id")
|
|
if sample_id in audits:
|
|
ordered.append(audits.pop(sample_id))
|
|
|
|
if audits:
|
|
raise ValueError(f"Found {len(audits)} audit IDs absent from the source sample")
|
|
if args.require_count is not None and len(ordered) != args.require_count:
|
|
raise ValueError(f"Expected {args.require_count} merged records, found {len(ordered)}")
|
|
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = args.output.with_suffix(args.output.suffix + ".tmp")
|
|
with temporary.open("w", encoding="utf-8") as handle:
|
|
for value in ordered:
|
|
handle.write(json.dumps(value, ensure_ascii=False, default=str) + "\n")
|
|
os.replace(temporary, args.output)
|
|
print(json.dumps({"output": str(args.output), "records": len(ordered)}, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|