scripts/add_missing_cas.py

#!/usr/bin/env python3
"""Add literature-derived CAS entries to the PINO registry."""
from __future__ import annotations

import argparse
import json
import logging
import sqlite3
import sys
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

from pino.registry import AromaRegistry

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
logger = logging.getLogger("add_missing_cas")


def collect_cas_from_literature(pino_root: Path) -> set[str]:
    """Gather all CAS numbers referenced by literature formulas."""
    all_cas: set[str] = set()
    formula_files = [
        pino_root / "data" / "literature_formulas.json",
        pino_root / "data" / "literature_formulas_poucher.jsonl",
        pino_root / "data" / "literature_formulas_appell_blocks.jsonl",
    ]
    for fpath in formula_files:
        if not fpath.exists():
            continue
        if fpath.suffix == ".jsonl":
            for line in fpath.read_text().strip().splitlines():
                rec = json.loads(line)
                for c in rec.get("components", []):
                    if c.get("type") == "pure" and c.get("cas"):
                        all_cas.add(c["cas"])
        else:
            data = json.loads(fpath.read_text())
            for rec in (data if isinstance(data, list) else [data]):
                for c in rec.get("components", []):
                    if c.get("type") == "pure" and c.get("cas"):
                        all_cas.add(c["cas"])
    return all_cas


def main() -> int:
    parser = argparse.ArgumentParser(description="Add missing CAS from literature to registry")
    parser.add_argument("--registry", default="src/pino/registry.db", help="Path to registry.db")
    parser.add_argument("--pino-root", default=".", help="Project root containing data/literature_formulas*")
    parser.add_argument("--limit", type=int, default=None, help="Only process first N missing CAS")
    parser.add_argument("--dry-run", action="store_true", help="Show missing CAS but do not modify")
    args = parser.parse_args()

    root = Path(args.pino_root).resolve()
    registry = AromaRegistry(args.registry)
    conn = sqlite3.connect(args.registry)
    cursor = conn.cursor()
    db_cas = {row[0] for row in cursor.execute("SELECT cas FROM aroma_chemicals WHERE cas IS NOT NULL")}
    conn.close()

    lit_cas = collect_cas_from_literature(root)
    missing = sorted(lit_cas - db_cas)
    logger.info("%d unique CAS in literature formulas; %d already in registry; %d missing", len(lit_cas), len(lit_cas & db_cas), len(missing))

    if args.dry_run:
        for cas in missing[: args.limit or len(missing)]:
            print(cas)
        return 0

    entries = [{"cas": cas} for cas in missing[: args.limit or len(missing)]]
    failed = registry.populate(entries, skip_failures=True)
    logger.info("Added %d/%d missing CAS; %d failed", len(entries) - len(failed), len(entries), len(failed))
    if failed:
        failed_path = root / "data" / "failed_literature_cas.json"
        failed_path.write_text(json.dumps(failed, indent=2))
        logger.info("Wrote failed CAS list to %s", failed_path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())