stimulus_synthesis/paths.py

"""Repo-local cache/path configuration.

Cache directory resolution order:
1. ``NEVO_CACHE_DIR`` environment variable (may be set via a ``.env`` file at the
   repository root),
2. otherwise ``<repo>/cache``.

``configure_cache()`` points HuggingFace/torch caches at that directory (using
``setdefault`` so any caller-provided env wins). No third-party dotenv dependency;
the ``.env`` parser here handles simple ``KEY=VALUE`` lines.
"""
from __future__ import annotations

import os
from pathlib import Path

_REPO_ROOT = Path(__file__).parent.parent


def _load_dotenv(root: Path) -> None:
    env_path = root / ".env"
    if not env_path.exists():
        return
    for line in env_path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        key = key.strip()
        value = value.strip().strip('"').strip("'")
        if key:
            os.environ.setdefault(key, value)


def get_cache_dir() -> Path:
    """Resolve and create the cache directory."""
    _load_dotenv(_REPO_ROOT)
    override = os.environ.get("NEVO_CACHE_DIR")
    path = Path(override).expanduser() if override else (_REPO_ROOT / "cache")
    path.mkdir(parents=True, exist_ok=True)
    return path


def _point_env_at(cache: Path) -> Path:
    cache.mkdir(parents=True, exist_ok=True)
    os.environ.setdefault("HF_HOME", str(cache / "huggingface"))
    os.environ.setdefault("HUGGINGFACE_HUB_CACHE", str(cache / "huggingface" / "hub"))
    os.environ.setdefault("TORCH_HOME", str(cache / "torch"))
    return cache


def configure_cache() -> Path | None:
    """Configure the HuggingFace/torch cache location.

    Priority:
    1. ``NEVO_CACHE_DIR`` (from environment or ``.env``) — use it.
    2. Otherwise, respect the system/user-default HuggingFace cache
       (``HF_HOME``/``HUGGINGFACE_HUB_CACHE`` if set, else ``~/.cache/huggingface``)
       and leave the environment untouched.
    3. Only if no default is resolvable (no usable home directory) fall back to
       ``<repo>/cache``.

    Returns the cache dir when this function set one, or ``None`` when the system
    default is left in place.
    """
    _load_dotenv(_REPO_ROOT)
    override = os.environ.get("NEVO_CACHE_DIR")
    if override:
        return _point_env_at(Path(override).expanduser())
    # Respect an already-configured / default HuggingFace cache.
    if os.environ.get("HF_HOME") or os.environ.get("HUGGINGFACE_HUB_CACHE"):
        return None
    home = os.path.expanduser("~")
    if home and home != "~" and os.path.isdir(home):
        return None  # HuggingFace will use its own default (~/.cache/huggingface)
    # No usable default: fall back to a repo-local cache.
    return _point_env_at(_REPO_ROOT / "cache")