src/slop_farmer/data/snapshot_source.py

from __future__ import annotations

from pathlib import Path

from slop_farmer.data.snapshot_materialize import materialize_hf_dataset_snapshot
from slop_farmer.data.snapshot_paths import (
    default_hf_materialize_dir,
    resolve_snapshot_dir_from_snapshots_root,
)


def resolve_snapshot_source_dir(
    *,
    snapshot_dir: Path | None,
    local_snapshots_root: Path,
    hf_repo_id: str | None,
    hf_revision: str | None,
    hf_materialize_dir: Path | None,
    hf_output_dir: Path | None = None,
) -> Path:
    if snapshot_dir is not None:
        return snapshot_dir.resolve()
    if hf_repo_id:
        output_dir = (hf_output_dir or local_snapshots_root.parent).resolve()
        return materialize_hf_dataset_snapshot(
            repo_id=hf_repo_id,
            local_dir=hf_materialize_dir
            or default_hf_materialize_dir(output_dir, hf_repo_id, hf_revision),
            revision=hf_revision,
        ).resolve()
    return resolve_snapshot_dir_from_snapshots_root(local_snapshots_root.resolve(), None)