src/fpgm/utils/io.py

"""Filesystem and HTTP download helpers shared across the data layer.

Every write to disk goes through :func:`atomic_write` or
:func:`download_with_resume`'s temp-file-then-rename pattern: a process killed
mid-download must never leave a truncated file sitting at the *final* path,
because that file's mere existence is what every later "already cached, skip
it" check relies on.
"""

from __future__ import annotations

import os
import tempfile
from pathlib import Path

import requests
from requests.adapters import HTTPAdapter
from tqdm import tqdm
from urllib3.util.retry import Retry

from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

_CHUNK_SIZE = 1024 * 1024  # 1 MiB


def ensure_dir(path: Path) -> Path:
    """Create ``path`` as a directory (parents included) if it doesn't exist."""
    path = Path(path)
    path.mkdir(parents=True, exist_ok=True)
    return path


def human_bytes(n: int | float) -> str:
    """Render a byte count as e.g. ``"3.5 GB"`` (decimal units, matching GCS/HF listings)."""
    n = float(n)
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1000.0:
            return f"{int(n)} {unit}" if unit == "B" else f"{n:.1f} {unit}"
        n /= 1000.0
    return f"{n:.1f} TB"


def atomic_write(path: Path, data: bytes) -> None:
    """Write ``data`` to ``path`` without ever exposing a partially-written file."""
    path = Path(path)
    ensure_dir(path.parent)
    fd, tmp_name = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.", suffix=".tmp")
    try:
        with os.fdopen(fd, "wb") as fh:
            fh.write(data)
        os.replace(tmp_name, path)
    except BaseException:
        Path(tmp_name).unlink(missing_ok=True)
        raise


def build_retrying_session(max_retries: int = 5, backoff_factor: float = 0.5) -> requests.Session:
    """A :class:`requests.Session` that retries transient failures with backoff."""
    session = requests.Session()
    retry = Retry(
        total=max_retries,
        backoff_factor=backoff_factor,
        status_forcelist=(429, 500, 502, 503, 504),
        allowed_methods=("GET", "HEAD"),
    )
    adapter = HTTPAdapter(max_retries=retry)
    session.mount("https://", adapter)
    session.mount("http://", adapter)
    return session


def remote_content_length(
    url: str, session: requests.Session, timeout: float = 30.0
) -> int | None:
    """``Content-Length`` of ``url`` via ``HEAD``, or ``None`` if unreported.

    Returns ``None`` when the response is content-encoded (gzip/br/deflate). In
    that case ``Content-Length`` describes the *compressed* transfer, while
    ``requests`` transparently decompresses the body before it reaches disk, so
    the two are not comparable. GitHub's raw CDN does exactly this for
    ``text/plain`` (our ``.dae`` and ``.urdf`` files, though not the
    ``application/octet-stream`` ``.stl``s), and comparing the two silently
    re-downloaded every text asset on every run.
    """
    resp = session.head(url, timeout=timeout, allow_redirects=True)
    resp.raise_for_status()

    encoding = resp.headers.get("Content-Encoding", "").strip().lower()
    if encoding and encoding != "identity":
        return None

    length = resp.headers.get("Content-Length")
    return int(length) if length is not None else None


def download_with_resume(
    url: str,
    dest: Path,
    *,
    session: requests.Session | None = None,
    timeout: float = 60.0,
    progress: bool = True,
) -> Path:
    """Stream ``url`` to ``dest``, skipping the download if it is already cached.

    "Resume" here is the pragmatic version that matters for this pipeline: a
    previously-completed download is recognised by comparing its size against
    the server's ``Content-Length`` and is never re-fetched, and an
    in-progress download is written to a sibling ``.part`` file so a crash
    never leaves a corrupt file at ``dest`` masquerading as a finished one.
    True byte-range resume is not implemented -- it isn't needed at these file
    sizes (mp4s ~5 MB, trajectory.h5 ~888 KB, metadata.json ~1.7 KB).

    Returns:
        ``dest``.
    """
    dest = Path(dest)
    session = session or build_retrying_session()
    ensure_dir(dest.parent)

    expected_size = remote_content_length(url, session, timeout=timeout)
    if dest.exists():
        actual_size = dest.stat().st_size
        if expected_size is not None and actual_size == expected_size:
            logger.debug("cached, skipping download: %s", dest)
            return dest
        if expected_size is None and actual_size > 0:
            # Size is unverifiable (compressed transfer or no Content-Length).
            # A complete file is still recognisable: partial downloads live in a
            # sibling ``.part`` and are only renamed once finished, so anything
            # sitting at ``dest`` with content is by construction complete.
            logger.debug("cached (size unverifiable), skipping download: %s", dest)
            return dest

    tmp_path = dest.with_name(dest.name + ".part")
    with session.get(url, stream=True, timeout=timeout) as resp:
        resp.raise_for_status()
        total = int(resp.headers.get("Content-Length", 0)) or expected_size or 0
        with (
            open(tmp_path, "wb") as fh,
            tqdm(
                total=total or None,
                unit="B",
                unit_scale=True,
                desc=dest.name,
                disable=not progress,
                leave=False,
            ) as bar,
        ):
            for chunk in resp.iter_content(chunk_size=_CHUNK_SIZE):
                if not chunk:
                    continue
                fh.write(chunk)
                bar.update(len(chunk))
    os.replace(tmp_path, dest)
    logger.info("downloaded %s (%s)", dest, human_bytes(dest.stat().st_size))
    return dest