Source code for vibesketch.inputs.registry

"""The registry of input writers, and the one call that runs them."""

from __future__ import annotations

from dataclasses import dataclass, field

from vibesketch.document.model import Document, JobSpec
from vibesketch.inputs.base import InputBundle, InputWriter, WriterRefusal

__all__ = [
    "Batch",
    "UnknownReference",
    "UnknownWriter",
    "get_writer",
    "register",
    "write_input",
    "write_inputs",
    "writers",
]

_WRITERS: dict[str, InputWriter] = {}




class UnknownWriter(LookupError):
    """No writer is registered under that name."""

    def __init__(self, name: str) -> None:
        self.name = name
        known = ", ".join(sorted(_WRITERS)) or "none"
        super().__init__(f"no input writer named {name!r}; available: {known}")





class UnknownReference(KeyError):
    """The document has no job or species with that id."""

    def __init__(self, kind: str, identifier: str) -> None:
        self.kind = kind
        self.identifier = identifier
        super().__init__(f"the document has no {kind} {identifier!r}")

    def __str__(self) -> str:
        return str(self.args[0])





def register(writer: InputWriter, *, replace: bool = False) -> None:
    """Make ``writer`` available under ``writer.name``."""
    if writer.name in _WRITERS and not replace:
        raise ValueError(f"an input writer named {writer.name!r} is already registered")
    _WRITERS[writer.name] = writer





def get_writer(name: str) -> InputWriter:
    """The writer registered as ``name``; raises :class:`UnknownWriter`."""
    try:
        return _WRITERS[name]
    except KeyError:
        raise UnknownWriter(name) from None





def writers() -> list[InputWriter]:
    """Every registered writer, by name."""
    return [_WRITERS[name] for name in sorted(_WRITERS)]





def write_input(
    document: Document,
    *,
    job: str | None = None,
    species: str | None = None,
    program: str | None = None,
    spec: JobSpec | None = None,
) -> InputBundle:
    """Write the input for one job of ``document``.

    Name a job by its id, or give a species id and a program (with an
    optional specification; the default is the geometry only). Raises
    :class:`~vibesketch.inputs.base.WriterRefusal` when the writer cannot
    express the request, :class:`UnknownWriter` for an unknown program, and
    :class:`UnknownReference` for an id the document does not have.
    """
    if job is not None:
        if species is not None or program is not None or spec is not None:
            raise TypeError("name a job, or a species and a program, not both")
        try:
            chosen = document.job_by_id(job)
        except KeyError:
            raise UnknownReference("job", job) from None
        # The model guarantees a job's species exists.
        return get_writer(chosen.program).write(document.species_by_id(chosen.species), chosen.spec)
    if species is None or program is None:
        raise TypeError("name a job, or a species and a program")
    try:
        target = document.species_by_id(species)
    except KeyError:
        raise UnknownReference("species", species) from None
    return get_writer(program).write(target, spec if spec is not None else JobSpec(task="geometry"))





@dataclass
class Batch:
    """Inputs for several species with one writer and one specification."""

    bundles: list[InputBundle] = field(default_factory=list)
    #: Per species id, why the writer refused it.
    refused: dict[str, list[str]] = field(default_factory=dict)





def write_inputs(
    document: Document,
    *,
    program: str,
    spec: JobSpec,
    species: list[str] | None = None,
) -> Batch:
    """The inputs for ``species`` (default: every species) with one writer.

    A species the writer refuses is listed in :attr:`Batch.refused` with the
    reasons, and the others are still written; a proton in a mechanism should
    not stop the rest of it. Raises :class:`UnknownWriter` for an unknown
    program and :class:`UnknownReference` for an unknown species id.
    """
    writer = get_writer(program)
    wanted = species if species is not None else [s.id for s in document.species]
    batch = Batch()
    for identifier in wanted:
        try:
            target = document.species_by_id(identifier)
        except KeyError:
            raise UnknownReference("species", identifier) from None
        try:
            batch.bundles.append(writer.write(target, spec))
        except WriterRefusal as refusal:
            batch.refused[identifier] = list(refusal.reasons)
    return batch