codekingpro/portable-devtools

1import re2from dataclasses import dataclass3from pathlib import Path4from typing import TYPE_CHECKING, Literal, TypedDict5 6from .. import constants7from ..file_download import repo_folder_name8from .sha import git_hash, sha_fileobj9 10 11if TYPE_CHECKING:12    from ..hf_api import RepoFile13 14# using fullmatch for clarity and strictness15_REGEX_COMMIT_HASH = re.compile(r"^[0-9a-f]{40}$")16 17 18# Typed structure describing a checksum mismatch19class Mismatch(TypedDict):20    path: str21    expected: str22    actual: str23    algorithm: str24 25 26HashAlgo = Literal["sha256", "git-sha1"]27 28 29@dataclass(frozen=True)30class FolderVerification:31    revision: str32    checked_count: int33    mismatches: list[Mismatch]34    missing_paths: list[str]35    extra_paths: list[str]36    verified_path: Path37 38 39def collect_local_files(root: Path) -> dict[str, Path]:40    """41    Return a mapping of repo-relative path -> absolute path for all files under `root`.42    """43    return {p.relative_to(root).as_posix(): p for p in root.rglob("*") if p.is_file()}44 45 46def _resolve_commit_hash_from_cache(storage_folder: Path, revision: str | None) -> str:47    """48    Resolve a commit hash from a cache repo folder and an optional revision.49    """50    if revision and _REGEX_COMMIT_HASH.fullmatch(revision):51        return revision52 53    refs_dir = storage_folder / "refs"54    snapshots_dir = storage_folder / "snapshots"55 56    if revision:57        ref_path = refs_dir / revision58        if ref_path.is_file():59            return ref_path.read_text(encoding="utf-8").strip()60        raise ValueError(f"Revision '{revision}' could not be resolved in cache (expected file '{ref_path}').")61 62    # No revision provided: try common defaults63    main_ref = refs_dir / "main"64    if main_ref.is_file():65        return main_ref.read_text(encoding="utf-8").strip()66 67    if not snapshots_dir.is_dir():68        raise ValueError(f"Cache repo is missing snapshots directory: {snapshots_dir}. Provide --revision explicitly.")69 70    candidates = [p.name for p in snapshots_dir.iterdir() if p.is_dir() and _REGEX_COMMIT_HASH.fullmatch(p.name)]71    if len(candidates) == 1:72        return candidates[0]73 74    raise ValueError(75        "Ambiguous cached revision: multiple snapshots found and no refs to disambiguate. Please pass --revision."76    )77 78 79def compute_file_hash(path: Path, algorithm: HashAlgo) -> str:80    """81    Compute the checksum of a local file using the requested algorithm.82    """83 84    with path.open("rb") as stream:85        if algorithm == "sha256":86            return sha_fileobj(stream).hex()87        if algorithm == "git-sha1":88            return git_hash(stream.read())89        raise ValueError(f"Unsupported hash algorithm: {algorithm}")90 91 92def verify_maps(93    *,94    remote_by_path: dict[str, "RepoFile"],95    local_by_path: dict[str, Path],96    revision: str,97    verified_path: Path,98) -> FolderVerification:99    """Compare remote entries and local files and return a verification result."""100    remote_paths = set(remote_by_path)101    local_paths = set(local_by_path)102 103    missing = sorted(remote_paths - local_paths)104    extra = sorted(local_paths - remote_paths)105    both = sorted(remote_paths & local_paths)106 107    mismatches: list[Mismatch] = []108 109    for rel_path in both:110        remote_entry = remote_by_path[rel_path]111        local_path = local_by_path[rel_path]112 113        lfs = getattr(remote_entry, "lfs", None)114        lfs_sha = getattr(lfs, "sha256", None) if lfs is not None else None115        if lfs_sha is None and isinstance(lfs, dict):116            lfs_sha = lfs.get("sha256")117        if lfs_sha:118            algorithm: HashAlgo = "sha256"119            expected = str(lfs_sha).lower()120        else:121            blob_id = remote_entry.blob_id  # type: ignore122            algorithm = "git-sha1"123            expected = str(blob_id).lower()124 125        actual = compute_file_hash(local_path, algorithm)126 127        if actual != expected:128            mismatches.append(Mismatch(path=rel_path, expected=expected, actual=actual, algorithm=algorithm))129 130    return FolderVerification(131        revision=revision,132        checked_count=len(both),133        mismatches=mismatches,134        missing_paths=missing,135        extra_paths=extra,136        verified_path=verified_path,137    )138 139 140def resolve_local_root(141    *,142    repo_id: str,143    repo_type: str,144    revision: str | None,145    cache_dir: Path | None,146    local_dir: Path | None,147) -> tuple[Path, str]:148    """149    Resolve the root directory to scan locally and the remote revision to verify.150    """151    if local_dir is not None:152        root = Path(local_dir).expanduser().resolve()153        if not root.is_dir():154            raise ValueError(f"Local directory does not exist or is not a directory: {root}")155        return root, (revision or constants.DEFAULT_REVISION)156 157    cache_root = Path(cache_dir or constants.HF_HUB_CACHE).expanduser().resolve()158    storage_folder = cache_root / repo_folder_name(repo_id=repo_id, repo_type=repo_type)159    if not storage_folder.exists():160        raise ValueError(161            f"Repo is not present in cache: {storage_folder}. Use 'hf download' first or pass --local-dir."162        )163    commit = _resolve_commit_hash_from_cache(storage_folder, revision)164    snapshot_dir = storage_folder / "snapshots" / commit165    if not snapshot_dir.is_dir():166        raise ValueError(f"Snapshot directory does not exist for revision '{commit}': {snapshot_dir}.")167    return snapshot_dir, commit168