scripts/verify_hub_dataset.py

#!/usr/bin/env python3
"""Confirm the Hub dataset has the catalog file the Space will download (no full download).

Uses Hub metadata APIs when available and falls back to a non-forced
``hf_hub_download`` for older ``huggingface_hub`` installs. For private
datasets, export HF_TOKEN first.

Example:

  HF_TOKEN=hf_... python scripts/verify_hub_dataset.py
  python scripts/verify_hub_dataset.py --repo-id ronheichman/kink-catalog-slim
"""
from __future__ import annotations

import argparse
import os
import sys


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--repo-id", default=os.environ.get("KINK_HF_DATASET_REPO", "ronheichman/kink-catalog-slim"))
    ap.add_argument("--filename", default=os.environ.get("KINK_HF_DATASET_FILENAME", "store_slim.db"))
    args = ap.parse_args()

    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")

    try:
        from huggingface_hub import HfApi

        api = HfApi(token=token)
        info = api.get_paths_info(
            args.repo_id,
            args.filename,
            repo_type="dataset",
            token=token,
        )[0]
    except Exception as e:
        try:
            from huggingface_hub import hf_hub_download

            path = hf_hub_download(
                args.repo_id,
                args.filename,
                repo_type="dataset",
                token=token,
                local_files_only=True,
            )
        except Exception:
            print(f"verify_hub_dataset: FAILED — {e}", file=sys.stderr)
            return 1
        from pathlib import Path

        p = Path(path)
        print(
            f"verify_hub_dataset: OK — {args.repo_id} / {args.filename} "
            f"(cached size {p.stat().st_size} bytes)",
            file=sys.stderr,
        )
        return 0

    print(
        f"verify_hub_dataset: OK — {args.repo_id} / {args.filename} "
        f"(size {getattr(info, 'size', '?')} bytes, blob {getattr(info, 'blob_id', '?')[:8]}…)",
        file=sys.stderr,
    )
    return 0


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