scripts/download_org_readmes.py

from __future__ import annotations

import json
import os
import re
import sys
import warnings
import shutil
from pathlib import Path

ROOT_DIR = Path(__file__).resolve().parents[1]
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))

from huggingface_hub import HfApi, hf_hub_download

from backend.spaces_index import INDEX_PATH, ORG_NAME, _extract_space_payload, _is_infra_repo, fetch_space_readme


DATA_DIR = ROOT_DIR / "data"
README_DIR = DATA_DIR / "org_readmes"
MANIFEST_PATH = DATA_DIR / "org_readmes_manifest.json"

os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")
warnings.filterwarnings(
    "ignore",
    message="`huggingface_hub` cache-system uses symlinks by default",
)


def _safe_slug(repo_id: str) -> str:
    return re.sub(r"[^a-zA-Z0-9._-]+", "__", repo_id.strip()) or "unknown-space"


def _download_readme(repo_id: str, api: HfApi) -> str:
    readme_text = fetch_space_readme(repo_id)
    if readme_text.strip():
        return readme_text

    try:
        repo_files = api.list_repo_files(repo_id=repo_id, repo_type="space")
    except Exception:
        return ""

    for repo_file in repo_files:
        normalized_name = Path(str(repo_file)).name.lower()
        if normalized_name != "readme.md":
            continue
        try:
            readme_path = hf_hub_download(repo_id=repo_id, repo_type="space", filename=str(repo_file))
            return Path(readme_path).read_text(encoding="utf-8", errors="ignore")
        except Exception:
            continue
    return ""


def _load_cached_index_targets() -> tuple[list[dict], dict]:
    if not INDEX_PATH.exists():
        return [], {}
    try:
        payload = json.loads(INDEX_PATH.read_text(encoding="utf-8"))
    except Exception:
        return [], {}
    if not isinstance(payload, dict):
        return [], {}
    spaces = payload.get("spaces", [])
    stats = payload.get("stats", {})
    if not isinstance(spaces, list):
        return [], {}

    targets: list[dict] = []
    for item in spaces:
        if not isinstance(item, dict):
            continue
        repo_id = str(item.get("id", "")).strip()
        if not repo_id:
            continue
        targets.append(
            {
                "repo_id": repo_id,
                "title": str(item.get("name", "")).strip(),
                "url": str(item.get("url", "")).strip(),
                "likes": int(item.get("likes", 0) or 0),
                "sdk": str(item.get("sdk", "")).strip(),
                "tags": list(item.get("tags", []) or []),
            }
        )
    return targets, stats if isinstance(stats, dict) else {}


def main() -> None:
    if README_DIR.exists():
        shutil.rmtree(README_DIR)
    README_DIR.mkdir(parents=True, exist_ok=True)
    api = HfApi()

    manifest: list[dict] = []
    readme_fetched = 0
    readme_missing = 0
    targets, cached_stats = _load_cached_index_targets()
    using_cached_index = bool(targets)

    if using_cached_index:
        listed = int(cached_stats.get("listed", len(targets)) or len(targets))
        private_skipped = int(cached_stats.get("private_skipped", 0) or 0)
        infra_skipped = int(cached_stats.get("infra_skipped", 0) or 0)
    else:
        try:
            raw_spaces = list(api.list_spaces(author=ORG_NAME, full=True))
        except TypeError:
            raw_spaces = list(api.list_spaces(author=ORG_NAME))

        listed = 0
        private_skipped = 0
        infra_skipped = 0
        targets = []
        for raw_space in raw_spaces:
            listed += 1
            payload = _extract_space_payload(raw_space)
            repo_id = payload["id"]
            status = str(payload.get("status", "")).strip().lower()
            if not repo_id:
                continue
            if payload.get("private") or payload.get("disabled") or payload.get("gated") or status == "private":
                private_skipped += 1
                continue
            if _is_infra_repo(repo_id, payload["raw_name"], payload["name"], payload["tags"]):
                infra_skipped += 1
                continue
            targets.append(
                {
                    "repo_id": repo_id,
                    "title": payload["name"],
                    "url": payload["url"],
                    "likes": int(payload.get("likes", 0) or 0),
                    "sdk": str(payload.get("sdk", "")).strip(),
                    "tags": list(payload.get("tags", []) or []),
                }
            )

    for index, target in enumerate(targets, start=1):
        repo_id = target["repo_id"]
        status = "unknown"

        readme_text = _download_readme(repo_id, api)
        saved = False
        relative_path = ""
        if readme_text.strip():
            folder = README_DIR / _safe_slug(repo_id)
            folder.mkdir(parents=True, exist_ok=True)
            readme_path = folder / "README.md"
            readme_path.write_text(readme_text, encoding="utf-8")
            relative_path = str(readme_path.relative_to(ROOT_DIR)).replace("\\", "/")
            saved = True
            readme_fetched += 1
        else:
            readme_missing += 1

        if index % 50 == 0:
            print(f"Processed {index}/{len(targets)} Spaces...")

        manifest.append(
            {
                "repo_id": repo_id,
                "title": target["title"],
                "url": target["url"],
                "status": status or "unknown",
                "likes": int(target.get("likes", 0) or 0),
                "sdk": str(target.get("sdk", "")).strip(),
                "tags": list(target.get("tags", []) or []),
                "readme_saved": saved,
                "readme_path": relative_path,
            }
        )

    manifest.sort(key=lambda item: (not item["readme_saved"], item["repo_id"].lower()))
    MANIFEST_PATH.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")

    saved_count = sum(1 for item in manifest if item["readme_saved"])
    missing_repo_ids = [item["repo_id"] for item in manifest if not item["readme_saved"]]
    if using_cached_index:
        print(f"Using cached index source: {INDEX_PATH}")
    print(f"Listed {listed} Spaces from {ORG_NAME}.")
    print(f"Private skipped: {private_skipped}")
    print(f"Infra skipped: {infra_skipped}")
    print(f"README fetched: {readme_fetched}")
    print(f"README missing: {readme_missing}")
    print(f"Indexed pool: {len(manifest)}")
    print(f"Saved {saved_count} README snapshots under {README_DIR}.")
    if missing_repo_ids:
        print("Missing README downloads:")
        for repo_id in missing_repo_ids:
            print(f"- {repo_id}")
    print(f"Manifest written to {MANIFEST_PATH}.")


if __name__ == "__main__":
    main()