"""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