stack/dev/scripts/add_carrier_to_zotero.py

"""Add 2010 and 2017 PFS carrier files to Zotero library.

Downloads from CMS, extracts, and creates Zotero items with proper
tags and storage-linked attachments.

Usage:
    uv run python dev/scripts/add_carrier_to_zotero.py
"""

import subprocess
import zipfile
from pathlib import Path

from conf import path as _conf_path
from zot.db import TYPE_MAP, Db, generate_key, now_iso

ZOTERO_DB = str(_conf_path("db.zotero"))
ZOTERO_STORAGE = str(_conf_path("storage.zotero"))

# Files already downloaded to /tmp
CARRIER_FILES = {
    2010: {
        "zip": "/tmp/pfs_carrier_dl/cy2010_carrier.zip",
        "extracted": "/tmp/pfs_carrier_dl/2010",
        "url": "https://www.cms.gov/medicare/medicare-fee-for-service-payment/physicianfeesched/downloads/cy2010qtr1carrierfiles3.zip",
        "title": "CY 2010 PFS Carrier — Cy2010 Carrier Files",
    },
    2017: {
        "zip": "/tmp/pfs_carrier_dl/cy2017_carrier.zip",
        "extracted": "/tmp/pfs_carrier_dl/2017",
        "url": "https://www.cms.gov/medicare/medicare-fee-for-service-payment/physicianfeesched/downloads/cy2017-carrierfiles.zip",
        "title": "CY 2017 PFS Carrier — Cy2017 Carrier Files",
    },
}


def add_carrier_year(db: Db, year: int, info: dict) -> None:
    """Add one carrier year to Zotero: parent item + file attachments."""
    now = now_iso()

    # Create parent webpage item
    parent_id = db.create_item(TYPE_MAP["webpage"], now=now)
    db.set_fields(
        parent_id,
        {
            "title": info["title"],
            "date": f"{year}-01-01",
            "url": info["url"],
            "accessDate": now,
            "websiteType": "Government Data Portal",
            "websiteTitle": "Centers for Medicare & Medicaid Services",
        },
    )
    db.sync_tags(parent_id, ["module:pfs", f"year:{year}"])

    print(f"Created parent item for {info['title']}")

    # Create attachments for each .TXT and .pdf file
    extracted_dir = Path(info["extracted"])
    att_count = 0

    for filepath in sorted(extracted_dir.iterdir()):
        if filepath.suffix.upper() not in (".TXT", ".PDF"):
            continue

        att_key = generate_key()
        storage_dir = Path(ZOTERO_STORAGE) / att_key
        subprocess.run(["sudo", "mkdir", "-p", str(storage_dir)], check=True)
        dest = storage_dir / filepath.name
        subprocess.run(["sudo", "cp", str(filepath), str(dest)], check=True)
        subprocess.run(
            ["sudo", "chown", "-R", "100999:100999", str(storage_dir)],
            check=True,
        )

        ext = filepath.suffix.upper()
        content_type = "text/plain" if ext == ".TXT" else "application/pdf"

        att_id = db.add_attachment(
            parent_id,
            key=att_key,
            content_type=content_type,
            path=f"storage:{filepath.name}",
        )
        db.set_field(att_id, "title", filepath.name)
        att_count += 1

    print(f"  Added {att_count} attachments for year {year}")


def main() -> None:
    for year, info in CARRIER_FILES.items():
        extracted = Path(info["extracted"])
        if not extracted.exists():
            print(f"Extracting {info['zip']}...")
            with zipfile.ZipFile(info["zip"]) as zf:
                zf.extractall(extracted)

        txt_count = len(list(extracted.glob("*.TXT")))
        print(f"Year {year}: {txt_count} .TXT files ready")

    with Db(ZOTERO_DB) as db:
        for year, info in sorted(CARRIER_FILES.items()):
            add_carrier_year(db, year, info)
        db.commit()
        print("\nDone. Committed to Zotero database.")


if __name__ == "__main__":
    main()