validation/fixture.py

"""Pinned public validation-fixture discovery, download, and verification."""

from __future__ import annotations

import hashlib
import os
import tempfile
import urllib.request
from dataclasses import dataclass
from pathlib import Path


@dataclass(frozen=True)
class PublicFixture:
    repository: str
    revision: str
    repository_path: str
    sha256: str
    size_bytes: int

    @property
    def url(self) -> str:
        return (
            f"https://huggingface.co/datasets/{self.repository}/resolve/"
            f"{self.revision}/{self.repository_path}?download=true"
        )


TESTING_FIXTURE = PublicFixture(
    repository="HWresearch/Delphes",
    revision="76a6c362bfba8e766ba7255a5c08a55257f78d0e",
    repository_path="testing/ttH_NLO_64.root",
    sha256="89d69eae9cd4d28a5d414d77bbc07185ea5b7de03481fdafe28899e2f3a09dec",
    size_bytes=15050,
)


def default_fixture_path(root: str | Path = ".") -> Path:
    return Path(root) / "data" / "fixtures" / "testing" / "ttH_NLO_64.root"


def file_sha256(path: str | Path) -> str:
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def verify_fixture(path: str | Path, fixture: PublicFixture = TESTING_FIXTURE) -> Path:
    """Verify size and SHA-256 before allowing a fixture into validation."""

    path = Path(path)
    if not path.is_file():
        raise FileNotFoundError(f"validation fixture is missing: {path}")
    actual_size = path.stat().st_size
    if actual_size != fixture.size_bytes:
        raise ValueError(
            f"fixture size mismatch for {path}: expected {fixture.size_bytes}, "
            f"got {actual_size}"
        )
    actual_sha256 = file_sha256(path)
    if actual_sha256 != fixture.sha256:
        raise ValueError(
            f"fixture SHA-256 mismatch for {path}: expected {fixture.sha256}, "
            f"got {actual_sha256}"
        )
    return path


def download_fixture(
    path: str | Path,
    fixture: PublicFixture = TESTING_FIXTURE,
    *,
    timeout: float = 60.0,
) -> Path:
    """Download a pinned fixture atomically and verify it before returning."""

    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary_name = tempfile.mkstemp(
        prefix=f".{path.name}.", suffix=".download", dir=path.parent
    )
    os.close(fd)
    temporary = Path(temporary_name)
    try:
        with urllib.request.urlopen(fixture.url, timeout=timeout) as response:
            with temporary.open("wb") as output:
                while chunk := response.read(1024 * 1024):
                    output.write(chunk)
        verify_fixture(temporary, fixture)
        os.replace(temporary, path)
    finally:
        temporary.unlink(missing_ok=True)
    return path


def ensure_fixture(
    path: str | Path | None = None,
    *,
    download: bool = True,
    fixture: PublicFixture = TESTING_FIXTURE,
) -> Path:
    """Return a verified local fixture, downloading it when necessary."""

    path = Path(path) if path is not None else default_fixture_path()
    if path.is_file():
        return verify_fixture(path, fixture)
    if not download:
        raise FileNotFoundError(f"validation fixture is missing: {path}")
    return download_fixture(path, fixture)