src/api/samples.py

"""Sample dataset discovery for the public example loader."""
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

_SUFFIX_EXTERIOR = "_exterior"
_SUFFIX_INTERIOR = "_interior"
_ALLOWED_EXT = (".step", ".stp")


@dataclass(frozen=True)
class SampleFile:
    filename: str
    path: Path
    size: int


@dataclass(frozen=True)
class SamplePair:
    id: str
    name: str
    exterior: SampleFile
    interior: SampleFile


def _pretty_name(stem: str) -> str:
    return stem.replace("_", " ").replace("-", " ").strip().title()


def _stem_for(path: Path, suffix: str) -> str | None:
    stem = path.stem
    return stem[: -len(suffix)] if stem.endswith(suffix) else None


def discover_samples(samples_dir: Path) -> list[SamplePair]:
    """Scan `samples_dir` for `{id}_exterior.(step|stp)` + `{id}_interior.(step|stp)` pairs."""
    if not samples_dir.exists() or not samples_dir.is_dir():
        return []

    by_id_ext: dict[str, Path] = {}
    by_id_int: dict[str, Path] = {}

    for file in sorted(samples_dir.iterdir()):
        if not file.is_file() or file.suffix.lower() not in _ALLOWED_EXT:
            continue
        if (sid := _stem_for(file, _SUFFIX_EXTERIOR)) is not None:
            by_id_ext[sid] = file
        elif (sid := _stem_for(file, _SUFFIX_INTERIOR)) is not None:
            by_id_int[sid] = file

    pairs: list[SamplePair] = []
    for sid in sorted(set(by_id_ext) & set(by_id_int)):
        ext_path = by_id_ext[sid]
        int_path = by_id_int[sid]
        pairs.append(
            SamplePair(
                id=sid,
                name=_pretty_name(sid),
                exterior=SampleFile(
                    filename=ext_path.name, path=ext_path, size=ext_path.stat().st_size
                ),
                interior=SampleFile(
                    filename=int_path.name, path=int_path, size=int_path.stat().st_size
                ),
            )
        )
    return pairs


def get_sample(samples_dir: Path, sample_id: str) -> SamplePair | None:
    for pair in discover_samples(samples_dir):
        if pair.id == sample_id:
            return pair
    return None