src/batch.py

import tempfile
import zipfile
from pathlib import Path

import pandas as pd

from . import pipeline
from .logger import log

SUPPORTED_EXTS = {".wav", ".mp3", ".ogg", ".flac", ".m4a"}


def _resolve_root(input_path: str) -> Path:
    p = Path(input_path)
    if p.is_dir():
        return p
    if p.suffix.lower() == ".zip":
        tmp = Path(tempfile.mkdtemp())
        with zipfile.ZipFile(p) as zf:
            zf.extractall(tmp)
        entries = list(tmp.iterdir())
        if len(entries) == 1 and entries[0].is_dir():
            return entries[0]
        return tmp
    raise ValueError(f"Expected a folder or .zip file, got: {input_path}")


def process_batch(input_path: str, on_progress=None) -> tuple[pd.DataFrame, list[dict]]:
    root = _resolve_root(input_path)
    errors: list[dict] = []

    csv_files = list(root.glob("*.csv"))
    manifest_names: set[str] | None = None
    if csv_files:
        manifest = pd.read_csv(csv_files[0])
        if "name" in manifest.columns:
            manifest_names = set(manifest["name"])
        else:
            errors.append(
                {"file": csv_files[0].name, "error": "manifest is missing the required 'name' column"}
            )

    audio_files = sorted(f for f in root.iterdir() if f.suffix.lower() in SUPPORTED_EXTS)
    if not audio_files:
        errors.append({"file": "", "error": "no supported audio files found in the batch"})

    if manifest_names is not None:
        found_names = {f.name for f in audio_files}
        for missing in sorted(manifest_names - found_names):
            errors.append({"file": missing, "error": "listed in manifest but not found in the uploaded batch"})

    rows = []
    for i, f in enumerate(audio_files, start=1):
        try:
            result = pipeline.analyze_one(str(f))
            rows.append({"name": f.name, **result.model_dump(mode="json")})
        except Exception as e:
            log.exception("failed to process %s: %s", f.name, e)
            errors.append({"file": f.name, "error": f"{type(e).__name__}: {e}"})
        if on_progress:
            on_progress(i, len(audio_files), f.name)

    return pd.DataFrame(rows), errors