Source code for gyoza.decorator.core

"""Simplified gyoza decorator for input/output model-driven function execution."""

from __future__ import annotations

from collections.abc import Callable
from functools import wraps
from typing import Any, TypeVar

from pydantic import BaseModel, ValidationError

F = TypeVar("F", bound=Callable[..., Any])
ModelType = type[BaseModel]


def _parse_output(result: Any, output_model: ModelType) -> BaseModel:
    """
    Convert a function's return value into the output model instance.

    Parameters
    ----------
    result : Any
        Raw return value from the decorated function.
    output_model : type[BaseModel]
        Pydantic model to validate and wrap the result.

    Returns
    -------
    BaseModel
        Validated output model instance.

    Raises
    ------
    ValueError
        If the result cannot be mapped to the output model.
    """
    if isinstance(result, output_model):
        return result

    if isinstance(result, BaseModel):
        return output_model.model_validate(result.model_dump())

    if isinstance(result, dict):
        return output_model(**result)

    if isinstance(result, tuple):
        fields = list(output_model.model_fields.keys())
        if len(fields) != len(result):
            msg = (
                f"Tuple result length {len(result)} does not match "
                f"output model field count {len(fields)}"
            )
            raise ValueError(msg)
        return output_model(**dict(zip(fields, result, strict=True)))

    fields = list(output_model.model_fields.keys())
    if len(fields) == 1:
        return output_model(**{fields[0]: result})

    msg = (
        f"Cannot convert result of type {type(result).__name__} "
        f"to {output_model.__name__}"
    )
    raise ValueError(msg)




class GyozaOp:
    """
    Bind a function to Pydantic input/output models.

    The decorated function keeps its original call signature and gains a
    ``run`` method that accepts a plain ``dict`` and returns a plain ``dict``,
    using the models for validation at both ends.

    Parameters
    ----------
    input_model : type[BaseModel]
        Pydantic model whose fields describe the function's inputs.
    output_model : type[BaseModel]
        Pydantic model whose fields describe the function's outputs.
    """

    def __init__(self, input_model: ModelType, output_model: ModelType) -> None:
        self.input_model = input_model
        self.output_model = output_model

    def __call__(self, func: F) -> F:
        """Wrap *func*, attaching ``input_model``, ``output_model`` and ``run``."""
        return self._wrap(func, self.input_model, self.output_model)

    @staticmethod
    def _wrap(
        func: F,
        input_model: ModelType,
        output_model: ModelType,
    ) -> F:
        @wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            return func(*args, **kwargs)

        def run(inputs: dict[str, Any]) -> dict[str, Any]:
            """
            Execute the function from a plain dict and return a plain dict.

            Parameters
            ----------
            inputs : dict[str, Any]
                Mapping whose keys match the fields of ``input_model``.

            Returns
            -------
            dict[str, Any]
                Mapping whose keys match the fields of ``output_model``.

            Raises
            ------
            ValueError
                If input validation or output conversion fails.
            """
            try:
                validated = input_model(**inputs)
            except ValidationError as exc:
                msg = f"Input validation error: {exc}"
                raise ValueError(msg) from exc

            result = func(
                **{k: getattr(validated, k) for k in type(validated).model_fields}
            )

            try:
                output = _parse_output(result, output_model)
            except (ValueError, ValidationError) as exc:
                msg = f"Output validation error: {exc}"
                raise ValueError(msg) from exc

            return output.model_dump()

        wrapper.run = run  # type: ignore[attr-defined]
        wrapper.input_model = input_model  # type: ignore[attr-defined]
        wrapper.output_model = output_model  # type: ignore[attr-defined]
        return wrapper  # type: ignore[return-value]





def gyoza_op(
    *,
    input_model: ModelType,
    output_model: ModelType,
) -> Callable[[F], F]:
    """
    Decorate a function with input/output model definitions.

    The decorated function retains its original call signature and gains:

    - ``run(inputs: dict) -> dict`` — validates inputs via ``input_model``,
      calls the function, and returns the result serialised via ``output_model``.
    - ``input_model`` / ``output_model`` — the attached model classes.

    Parameters
    ----------
    input_model : type[BaseModel]
        Pydantic model whose fields describe the function's inputs.
    output_model : type[BaseModel]
        Pydantic model whose fields describe the function's outputs.

    Returns
    -------
    Callable[[F], F]
        Decorator that wraps the target function.
    """
    return GyozaOp(input_model=input_model, output_model=output_model)