product-54354138-1/tools/envelope.py

"""Status vocabulary, exit codes and the JSON response envelope.

Every invocation of ``python -m tools.example`` writes exactly one JSON
document to stdout.  The document always carries a ``status`` value from the
fixed vocabulary below and the process exit code is derived from that value, so
a calling AI agent can tell a successful search apart from "nothing matched",
a broken request and an unusable source without parsing prose.
"""

from __future__ import annotations

import json
import sys
from dataclasses import dataclass
from typing import Any, Iterable, Mapping, TextIO

from .errors import (
    RequestError,
    SearchRequestError,
    SourceAccessError,
    ToolkitError,
    UnknownSourceError,
)

__all__ = [
    "EXIT_ACCESS_ERROR",
    "EXIT_INVALID_REQUEST",
    "EXIT_NO_FORMAT",
    "EXIT_NO_SIZE",
    "EXIT_OK",
    "EXIT_PARTIAL_DOWNLOAD",
    "Response",
    "ResponseBuilder",
    "STATUS_ACCESS_ERROR",
    "STATUS_DRY_RUN",
    "STATUS_EXIT_CODES",
    "STATUS_INVALID_REQUEST",
    "STATUS_NO_FORMAT",
    "STATUS_NO_SIZE",
    "STATUS_OK",
    "STATUS_PARTIAL_DOWNLOAD",
    "STATUS_VOCABULARY",
    "error_fragment",
    "exit_code_for",
    "is_success",
    "status_for_error",
]

#: Successful search/download.
STATUS_OK = "ok"
#: The request was validated and echoed without contacting a source.
STATUS_DRY_RUN = "dry_run"
#: The request itself is unusable (bad flags, bad document, unknown source or format).
STATUS_INVALID_REQUEST = "invalid_request"
#: The source answered but offers nothing in the requested format.
STATUS_NO_FORMAT = "no_format"
#: The source answered but offers nothing in the requested size.
STATUS_NO_SIZE = "no_size"
#: The source could not be reached, refused to answer or returned unusable data.
STATUS_ACCESS_ERROR = "access_error"
#: Some downloads succeeded and some failed.
STATUS_PARTIAL_DOWNLOAD = "partial_download"

#: The single source of truth for the status -> exit code mapping.
STATUS_EXIT_CODES: dict[str, int] = {
    STATUS_OK: 0,
    STATUS_DRY_RUN: 0,
    STATUS_INVALID_REQUEST: 2,
    STATUS_NO_FORMAT: 3,
    STATUS_NO_SIZE: 4,
    STATUS_ACCESS_ERROR: 5,
    STATUS_PARTIAL_DOWNLOAD: 6,
}

#: Stable ordering used in help text and documentation.
STATUS_VOCABULARY: tuple[str, ...] = tuple(STATUS_EXIT_CODES)

EXIT_OK = STATUS_EXIT_CODES[STATUS_OK]
EXIT_INVALID_REQUEST = STATUS_EXIT_CODES[STATUS_INVALID_REQUEST]
EXIT_NO_FORMAT = STATUS_EXIT_CODES[STATUS_NO_FORMAT]
EXIT_NO_SIZE = STATUS_EXIT_CODES[STATUS_NO_SIZE]
EXIT_ACCESS_ERROR = STATUS_EXIT_CODES[STATUS_ACCESS_ERROR]
EXIT_PARTIAL_DOWNLOAD = STATUS_EXIT_CODES[STATUS_PARTIAL_DOWNLOAD]


def exit_code_for(status: str) -> int:
    """Return the process exit code for ``status``.

    Raises :class:`ValueError` for anything outside the fixed vocabulary, so a
    typo can never silently produce an unknown exit code.
    """

    try:
        return STATUS_EXIT_CODES[status]
    except KeyError:
        raise ValueError(
            f"unknown status {status!r}; expected one of: {', '.join(STATUS_VOCABULARY)}"
        ) from None


def is_success(status: str) -> bool:
    """True when ``status`` maps onto a zero exit code."""

    return exit_code_for(status) == 0


def error_fragment(
    code: str,
    message: str,
    *,
    details: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
    """Build the machine-readable ``error`` fragment of a response."""

    payload: dict[str, Any] = {"code": str(code), "message": str(message)}
    if details:
        payload["details"] = dict(details)
    return payload


def status_for_error(error: BaseException) -> str:
    """Map a toolkit error onto the fixed status vocabulary."""

    if isinstance(error, (RequestError, UnknownSourceError, SearchRequestError)):
        return STATUS_INVALID_REQUEST
    if isinstance(error, SourceAccessError):
        return STATUS_ACCESS_ERROR
    if isinstance(error, ToolkitError):
        return STATUS_ACCESS_ERROR
    return STATUS_ACCESS_ERROR


@dataclass(frozen=True)
class Response:
    """One complete machine-readable answer of the toolkit."""

    status: str
    source: str | None = None
    request: Mapping[str, Any] | None = None
    items: tuple[Mapping[str, Any], ...] = ()
    error: Mapping[str, Any] | None = None

    def __post_init__(self) -> None:
        if self.status not in STATUS_EXIT_CODES:
            raise ValueError(
                f"unknown status {self.status!r}; expected one of: "
                f"{', '.join(STATUS_VOCABULARY)}"
            )
        if self.error is not None and "code" not in self.error:
            raise ValueError("error fragment must carry a 'code'")

    @property
    def exit_code(self) -> int:
        """Process exit code implied by :attr:`status`."""

        return exit_code_for(self.status)

    @property
    def succeeded(self) -> bool:
        """True when the status is one of the zero-exit statuses."""

        return self.exit_code == 0

    def to_dict(self) -> dict[str, Any]:
        """Return the JSON-ready document with a fixed set of top-level keys."""

        return {
            "status": self.status,
            "source": self.source,
            "request": dict(self.request) if self.request is not None else None,
            "items": [dict(item) for item in self.items],
            "error": dict(self.error) if self.error is not None else None,
        }

    def to_json(self) -> str:
        """Serialise the response as one deterministic JSON line."""

        return json.dumps(self.to_dict(), sort_keys=True, ensure_ascii=True)

    def emit(self, stream: TextIO | None = None) -> None:
        """Write exactly one JSON document to ``stream`` (stdout by default)."""

        print(self.to_json(), file=stream if stream is not None else sys.stdout)


class ResponseBuilder:
    """Build :class:`Response` objects sharing one source/request context."""

    def __init__(
        self,
        *,
        source: str | None = None,
        request: Mapping[str, Any] | None = None,
    ) -> None:
        self.source = source
        self.request = request

    def with_context(
        self,
        *,
        source: str | None = None,
        request: Mapping[str, Any] | None = None,
    ) -> "ResponseBuilder":
        """Return a copy carrying the supplied context."""

        return ResponseBuilder(
            source=self.source if source is None else source,
            request=self.request if request is None else request,
        )

    def build(
        self,
        status: str,
        *,
        items: Iterable[Mapping[str, Any]] = (),
        error: Mapping[str, Any] | None = None,
    ) -> Response:
        """Assemble a response for an explicit ``status``."""

        return Response(
            status=status,
            source=self.source,
            request=self.request,
            items=tuple(items),
            error=error,
        )

    # -- convenience constructors --------------------------------------
    def ok(self, *, items: Iterable[Mapping[str, Any]] = ()) -> Response:
        return self.build(STATUS_OK, items=items)

    def dry_run(self, *, items: Iterable[Mapping[str, Any]] = ()) -> Response:
        return self.build(STATUS_DRY_RUN, items=items)

    def invalid_request(self, *, error: Mapping[str, Any] | None = None) -> Response:
        return self.build(STATUS_INVALID_REQUEST, error=error)

    def no_format(self, *, error: Mapping[str, Any] | None = None) -> Response:
        return self.build(STATUS_NO_FORMAT, error=error)

    def no_size(self, *, error: Mapping[str, Any] | None = None) -> Response:
        return self.build(STATUS_NO_SIZE, error=error)

    def access_error(
        self,
        *,
        items: Iterable[Mapping[str, Any]] = (),
        error: Mapping[str, Any] | None = None,
    ) -> Response:
        return self.build(STATUS_ACCESS_ERROR, items=items, error=error)

    def partial_download(
        self,
        *,
        items: Iterable[Mapping[str, Any]] = (),
        error: Mapping[str, Any] | None = None,
    ) -> Response:
        return self.build(STATUS_PARTIAL_DOWNLOAD, items=items, error=error)

    def from_error(self, error: ToolkitError) -> Response:
        """Build a response for ``error`` using :func:`status_for_error`."""

        return self.build(status_for_error(error), error=error.to_dict())