from __future__ import annotations from dataclasses import dataclass from datetime import datetime from pathlib import Path from typing import Any import pandas as pd from huggingface_hub import HfApi, hf_hub_download from huggingface_hub.utils import EntryNotFoundError, HfHubHTTPError from config import AppConfig README_CANDIDATES = ("README.md", "readme.md") @dataclass(frozen=True) class RepoRevision: repo_id: str revision: str committed_at: datetime | None class HubClient: def __init__(self, config: AppConfig): self.config = config self.api = HfApi(token=config.hf_token) def discover_dataset_repos(self) -> list[str]: repos = set(self.config.repo_allowlist) if self.config.discovery_enabled: try: datasets = self.api.list_datasets( author=self.config.hf_owner, search=self.config.repo_id_contains or None, limit=max(self.config.max_repos * 3, self.config.max_repos), full=False, ) for dataset in datasets: repo_id = getattr(dataset, "id", None) if not repo_id: continue if self._repo_matches_filters(repo_id): repos.add(repo_id) if len(repos) >= self.config.max_repos: break except HfHubHTTPError: pass filtered = [repo for repo in repos if self._repo_matches_filters(repo)] return sorted(filtered)[: self.config.max_repos] def list_dataset_commits(self, repo_id: str, max_commits: int) -> list[RepoRevision]: revisions: list[RepoRevision] = [] try: commits = self.api.list_repo_commits(repo_id=repo_id, repo_type="dataset") for commit in commits[:max_commits]: revision = getattr(commit, "commit_id", None) or getattr(commit, "oid", None) if not revision: continue committed_at = _coerce_datetime(getattr(commit, "created_at", None) or getattr(commit, "date", None)) revisions.append( RepoRevision( repo_id=repo_id, revision=revision, committed_at=committed_at, ) ) except HfHubHTTPError: revisions = [] if revisions: return revisions try: info = self.api.dataset_info(repo_id=repo_id) fallback_revision = getattr(info, "sha", None) or "main" fallback_timestamp = _coerce_datetime(getattr(info, "lastModified", None)) except HfHubHTTPError: fallback_revision = "main" fallback_timestamp = None return [RepoRevision(repo_id=repo_id, revision=fallback_revision, committed_at=fallback_timestamp)] def fetch_dataset_readme(self, repo_id: str, revision: str | None = None) -> str | None: for filename in README_CANDIDATES: try: path = hf_hub_download( repo_id=repo_id, filename=filename, repo_type="dataset", revision=revision, token=self.config.hf_token, ) return Path(path).read_text(encoding="utf-8", errors="replace") except EntryNotFoundError: continue except HfHubHTTPError: continue return None def _repo_matches_filters(self, repo_id: str) -> bool: lowered_repo_id = repo_id.lower() if self.config.repo_id_contains and self.config.repo_id_contains.lower() not in lowered_repo_id: return False if self.config.repo_prefixes and not any(repo_id.startswith(prefix) for prefix in self.config.repo_prefixes): return False return True def _coerce_datetime(value: Any) -> datetime | None: if value is None: return None timestamp = pd.to_datetime(value, errors="coerce", utc=True) if pd.isna(timestamp): return None return timestamp.to_pydatetime()